Hi all, I couldn't use random.PRNGKey to generate random seed. Wondering anyone has similar issue before and figure this out. Here is my current config: jax-metal==0.0.3, jaxlib==0.4.10, jax==0.4.11. I am using Apple M1 pro.

Hi, this is the 3rd time I'm trying to post this on the forum, apple moderators ignoring it. I'm a deep learning expert with a specialization of image processing. I want to know why I have hundreds of AI models on my Mac that are indexing everything on my computer while it is idle, using programs like neuralhash that I can't find any information about. I can understand if they are being used to enhance the user experience on Spotlight, Siri, Photos, and other applications, but I couldn't find the necessary information on the web. Usually, (spyware) software like this uses them to classify files in an X/Y coordinate system. This feels like a more advanced version of stuxnet. find / -type f -name "*.weights" > ai_models.txt find / -type f -name "*labels*.txt" > ai_model_labels.txt Some of the classes from the files; file_name: SCL_v0.3.1_9c7zcipfrc_558001-labels-v3.txt document_boarding_pass document_check_or_checkbook document_currency_or_bill document_driving_license document_office_badge document_passport document_receipt document_social_security_number hier_curation hier_document hier_negative curation_meme file_name: SceneNet5_detection_labels-v8d.txt CVML_UNKNOWN_999999 aircraft automobile bicycle bird bottle bus canine consumer_electronics feline fruit furniture headgear kite fish computer_monitor motorcycle musical_instrument document people food sign watersport train ungulates watercraft flower appliance sports_equipment tool

Hi I have only recently started working on ML on my Mac M1 Pro, previously I was working on a Windows platform. I am having difficulties getting my machine set up right so that its ready for the super fast training I was hoping for when I got it. Please help me with this and let me know if and where I am going wrong. So, I tried a custom dataset training using Yolov8 model. I want to train for a 100 epochs. Now the same dataset and hyperparameters take about 2.5 hours on a T4 GPU on Google Colab, whereas I was only at around 60 epochs after 24 hours on my M1 pro. I have home brew, miniconda, pytorch nightly for mac installed and set the device to mps when training the YOLO model. I feel that this is reaaallly slow. What should I be doing right? Thank you Lakshmi

followed instruction in https://developer.apple.com/metal/jax/ I got Successfully installed importlib-metadata-7.1.0 jax-0.4.28 jax-metal-0.0.7 jaxlib-0.4.28 opt-einsum-3.3.0 scipy-1.13.0 six-1.16.0 zipp-3.18.2 but the test failed python -c 'import jax; print(jax.numpy.arange(10))' Traceback (most recent call last): File "", line 1, in File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/init.py", line 37, in import jax.core as _core File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/core.py", line 18, in from jax._src.core import ( File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/_src/core.py", line 39, in from jax._src import dtypes File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/_src/dtypes.py", line 33, in from jax._src import config File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/_src/config.py", line 27, in from jax._src import lib File "/Users/erivas/jax-metal/lib/python3.9/site-packages/jax/_src/lib/init.py", line 84, in cpu_feature_guard.check_cpu_features() RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

Cannot assign a device for operation encoder/down1/downs_0/conv1/weight/Initializer/random_uniform/RandomUniform: Could not satisfy explicit device specification '' because the node {{colocation_node encoder/down1/downs_0/conv1/weight/Initializer/random_uniform/RandomUniform}} was colocated with a group of nodes that required incompatible device '/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. Colocation Debug Info: Colocation group had the following types and supported devices: Root Member(assigned_device_name_index_=-1 requested_device_name_='/device:GPU:0' assigned_device_name_='' resource_device_name_='/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[] Identity: GPU CPU Mul: GPU CPU AddV2: GPU CPU Sub: GPU CPU RandomUniform: GPU CPU Assign: CPU VariableV2: GPU CPU Const: GPU CPU

Hi all, I'm having trouble even getting jax-metal latest version to install on my M1 MacBook Pro. In a clean conda environment, I pip install jax-metal and get In [1]: import jax; print(jax.numpy.arange(10)) Platform 'METAL' is experimental and not all JAX functionality may be correctly supported! --------------------------------------------------------------------------- XlaRuntimeError Traceback (most recent call last) [... skipping hidden 1 frame] File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/xla_bridge.py:977, in _init_backend(platform) 976 logger.debug("Initializing backend '%s'", platform) --> 977 backend = registration.factory() 978 # TODO(skye): consider raising more descriptive errors directly from backend 979 # factories instead of returning None. File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/xla_bridge.py:666, in register_plugin.<locals>.factory() 665 if not xla_client.pjrt_plugin_initialized(plugin_name): --> 666 xla_client.initialize_pjrt_plugin(plugin_name) 667 updated_options = {} File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jaxlib/xla_client.py:176, in initialize_pjrt_plugin(plugin_name) 169 """Initializes a PJRT plugin. 170 171 The plugin needs to be loaded first (through load_pjrt_plugin_dynamically or (...) 174 plugin_name: the name of the PJRT plugin. 175 """ --> 176 _xla.initialize_pjrt_plugin(plugin_name) XlaRuntimeError: INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.47) and framework PJRT API version 0.51). During handling of the above exception, another exception occurred: RuntimeError Traceback (most recent call last) Cell In[1], line 1 ----> 1 import jax; print(jax.numpy.arange(10)) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:2952, in arange(start, stop, step, dtype) 2950 ceil_ = ufuncs.ceil if isinstance(start, core.Tracer) else np.ceil 2951 start = ceil_(start).astype(int) # type: ignore -> 2952 return lax.iota(dtype, start) 2953 else: 2954 if step is None and start == 0 and stop is not None: File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/lax/lax.py:1282, in iota(dtype, size) 1277 def iota(dtype: DTypeLike, size: int) -> Array: 1278 """Wraps XLA's `Iota 1279 <https://www.tensorflow.org/xla/operation_semantics#iota>`_ 1280 operator. 1281 """ -> 1282 return broadcasted_iota(dtype, (size,), 0) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/lax/lax.py:1292, in broadcasted_iota(dtype, shape, dimension) 1289 static_shape = [None if isinstance(d, core.Tracer) else d for d in shape] 1290 dimension = core.concrete_or_error( 1291 int, dimension, "dimension argument of lax.broadcasted_iota") -> 1292 return iota_p.bind(*dynamic_shape, dtype=dtype, shape=tuple(static_shape), 1293 dimension=dimension) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/core.py:387, in Primitive.bind(self, *args, **params) 384 def bind(self, *args, **params): 385 assert (not config.enable_checks.value or 386 all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args --> 387 return self.bind_with_trace(find_top_trace(args), args, params) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/core.py:391, in Primitive.bind_with_trace(self, trace, args, params) 389 def bind_with_trace(self, trace, args, params): 390 with pop_level(trace.level): --> 391 out = trace.process_primitive(self, map(trace.full_raise, args), params) 392 return map(full_lower, out) if self.multiple_results else full_lower(out) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/core.py:879, in EvalTrace.process_primitive(self, primitive, tracers, params) 877 return call_impl_with_key_reuse_checks(primitive, primitive.impl, *tracers, **params) 878 else: --> 879 return primitive.impl(*tracers, **params) File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/dispatch.py:86, in apply_primitive(prim, *args, **params) 84 prev = lib.jax_jit.swap_thread_local_state_disable_jit(False) 85 try: ---> 86 outs = fun(*args) 87 finally: 88 lib.jax_jit.swap_thread_local_state_disable_jit(prev) [... skipping hidden 17 frame] File ~/opt/anaconda3/envs/metal/lib/python3.11/site-packages/jax/_src/xla_bridge.py:902, in backends() 900 else: 901 err_msg += " (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)" --> 902 raise RuntimeError(err_msg) 904 assert _default_backend is not None 905 if not config.jax_platforms.value: RuntimeError: Unable to initialize backend 'METAL': INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.47) and framework PJRT API version 0.51). (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.) jax.__version__ is 0.4.27.

Hello there, We currently have a crash in prod when executing the following line: let classificationRequest = try SNClassifySoundRequest(classifierIdentifier: .version1) It appears to only happen on iOS 17+ and only when regaining audio focus after an interruption in a background state. We are aware this call probably fails because it is happening from a background state - however - I would expect then that the SNClassifySoundRequest throws some kind of error since it is already an initializer that throws. If it is possible for the call to fail under certain circumstances, then could SNMLModelFactory throw an error instead of using try! ? Full trace below: SoundAnalysis/SNMLModelFactory.swift:112: Fatal error: 'try!' expression unexpectedly raised an error: Error Domain=com.apple.CoreML Code=0 "Failed to build the model execution plan using a model architecture file '/System/Library/Frameworks/SoundAnalysis.framework/SNSoundClassifierVersion1Model.mlmodelc/model1/model.espresso.net' with error code: -1." UserInfo={NSLocalizedDescription=Failed to build the model execution plan using a model architecture file '/System/Library/Frameworks/SoundAnalysis.framework/SNSoundClassifierVersion1Model.mlmodelc/model1/model.espresso.net' with error code: -1.}

Hi everyone ! I'm getting random crashes when I'm using the Speech Recognizer functionality in my app. This is an old bug (for 8 years on Apple Forums) and I will really appreciate if anyone from Apple will be able to find a fix for this crashes. Can anyone also help me please to understand what could I do to keep the Speech Recognizer functionality still available in my app, but to avoid this crashes (if there is any other native library available or a CocoaPod library). Here is my code and also the crash log for it. Code: func startRecording() { startStopRecordBtn.setImage(UIImage(#imageLiteral(resourceName: "microphone_off")), for: .normal) if UserDefaults.standard.bool(forKey: Constants.darkTheme) { commentTextView.textColor = .white } else { commentTextView.textColor = .black } commentTextView.isUserInteractionEnabled = false recordingLabel.text = Constants.recording if recognitionTask != nil { recognitionTask?.cancel() recognitionTask = nil } let audioSession = AVAudioSession.sharedInstance() do { try audioSession.setCategory(AVAudioSession.Category.record) try audioSession.setMode(AVAudioSession.Mode.measurement) try audioSession.setActive(true, options: .notifyOthersOnDeactivation) } catch { showAlertWithTitle(message: Constants.error) } recognitionRequest = SFSpeechAudioBufferRecognitionRequest() let inputNode = audioEngine.inputNode guard let recognitionRequest = recognitionRequest else { fatalError(Constants.error) } recognitionRequest.shouldReportPartialResults = true recognitionTask = speechRecognizer?.recognitionTask(with: recognitionRequest, resultHandler: { (result, error) in var isFinal = false if result != nil { self.commentTextView.text = result?.bestTranscription.formattedString isFinal = (result?.isFinal)! } if error != nil || isFinal { self.audioEngine.stop() inputNode.removeTap(onBus: 0) self.recognitionRequest = nil self.recognitionTask = nil self.startStopRecordBtn.isEnabled = true } }) let recordingFormat = inputNode.outputFormat(forBus: 0) inputNode.installTap(onBus: 0, bufferSize: 1024, format: recordingFormat) {[weak self] (buffer: AVAudioPCMBuffer, when: AVAudioTime) in // CRASH HERE self?.recognitionRequest?.append(buffer) } audioEngine.prepare() do { try audioEngine.start() } catch { showAlertWithTitle(message: Constants.error) } } Here is the crash log: Thanks for very much for reading this !

Regardless of the installation version combinations of tensorflow & metal (2.14, 2.15, 2.16), I find a metal/non-metal incompatibility for some layer types. For the GRU layer, for example, metal-trained weights (model.save_weights()/load_weights()) are not compatible with inference using the CPU. That is, train a model using metal, run inference using metal, then run inference again after uninstalling metal, and the results differ -- sometimes a night and day difference. This essentially eliminates the usefulness of tensorflow-metal for me. From my limited testing, models using other, simple combinations of layer types including Dense and LSTM do not show this problem. Just the GRU. And by "testing" I mean really simple models, like one GRU layer. Apple Framework Metal Team: You are doing very useful work, and I kindly ask, please address this bug :)

I noticed from the system requirements, TensorFlow only seems to support Python. Are there any plans to add JavaScript as TensorFlow has JS support? Thank you for your time...

I using a Macbook pro with an m2 pro chip. I was trying to work with TensorFlow but I encountered an illegal hardware instruction error. To resolve it I initiated the installation of a metal plugin which is throwing the following error. or semicolon (after version specifier) awscli>=1.16.100boto3>=1.9.100 ~~~~~~~~~~~^ Unable to locate awscli [end of output]

Hello! I am developing an app that leverages Apple's 2D pose estimation model and I would love to speak with someone about if my mobile app should leverage Apple's 3D pose estimation model. Also, I would love to know if Apple considers adding more points on the body as this would be incredibly helpful. Or if it is possible for me to train the model to add more body points. Thanks so much and please let me know if anyone is available to discuss.

When fitting a CNN model, every second Epoch takes zero seconds and with OUT_OF_RANGE warnings. Im using structured folders of categorical images for training and validation. Here is the warning message that occurs after every second Epoch. The fitting looks like this... 37/37 ━━━━━━━━━━━━━━━━━━━━ 14s 337ms/step - accuracy: 0.5255 - loss: 1.0819 - val_accuracy: 0.2578 - val_loss: 2.4472 Epoch 4/20 37/37 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step - accuracy: 0.5312 - loss: 1.1106 - val_accuracy: 0.1250 - val_loss: 3.0711 Epoch 5/20 2024-04-19 09:22:51.673909: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence [[{{node IteratorGetNext}}]] 2024-04-19 09:22:51.673928: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence [[{{node IteratorGetNext}}]] [[IteratorGetNext/_59]] 2024-04-19 09:22:51.673940: I tensorflow/core/framework/local_rendezvous.cc:422] Local rendezvous recv item cancelled. Key hash: 10431687783238222105 2024-04-19 09:22:51.673944: I tensorflow/core/framework/local_rendezvous.cc:422] Local rendezvous recv item cancelled. Key hash: 17360824274615977051 2024-04-19 09:22:51.673955: I tensorflow/core/framework/local_rendezvous.cc:422] Local rendezvous recv item cancelled. Key hash: 10732905483452597729 My setup is.. Tensor Flow Version: 2.16.1 Python 3.9.19 (main, Mar 21 2024, 12:07:41) [Clang 14.0.6 ] Pandas 2.2.2 Scikit-Learn 1.4.2 GPU is available My generator is.. train_generator = datagen.flow_from_directory( scalp_dir_train, # directory target_size=(256, 256),# all images found will be resized batch_size=32, class_mode='categorical' #subset='training' # Specify the subset as training ) n_samples = train_generator.samples # gets the number of samples validation_generator = datagen.flow_from_directory( scalp_dir_test, # directory path target_size=(256, 256), batch_size=32, class_mode='categorical' #subset='validation' # Specifying the subset as validation Here is my model. early_stopping_monitor = EarlyStopping(patience = 10,restore_best_weights=True) from tensorflow.keras.optimizers import Adam from tensorflow.keras.optimizers import SGD optimizer = Adam(learning_rate=0.01) model = Sequential() model.add(Conv2D(128, (3, 3), activation='relu',padding='same', input_shape=(256, 256, 3))) model.add(BatchNormalization()) model.add(MaxPooling2D((2, 2))) model.add(Dropout(0.3)) model.add(Conv2D(64, (3, 3),padding='same', activation='relu')) model.add(BatchNormalization()) model.add(MaxPooling2D((2, 2))) model.add(Dropout(0.3)) model.add(Flatten()) model.add(Dense(512, activation='relu')) model.add(BatchNormalization()) model.add(Dropout(0.4)) model.add(Dense(256, activation='relu')) model.add(BatchNormalization()) model.add(Dropout(0.3)) model.add(Dense(4, activation='softmax')) # Defined by the number of classes model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy']) Here is the fit... history=model.fit( train_generator, steps_per_epoch=37, epochs=20, validation_data=validation_generator, validation_steps=12, callbacks=[early_stopping_monitor] #verbose=2 )

Hi, I just noticed that using the jax.numpy.insert() function returns an incorrect result (zero-padding the array) when compiled with jax.jit. When not jitted, the results are correct Config: M1 Pro Macbook Pro 2021 python 3.12.3 ; jax-metal 0.0.6 ; jax 0.4.26 ; jaxlib 0.4.23 MWE: import jax import jax.numpy as jnp x = jnp.arange(20).reshape(5, 4) print(f"{x=}\n") def return_arr_with_ins(arr, ins): return jnp.insert(arr, 2, ins, axis=1) x2 = return_arr_with_ins(x, 99) print(f"{x2=}\n") return_arr_with_ins_jit = jax.jit(return_arr_with_ins) x3 = return_arr_with_ins_jit(x, 99) print(f"{x3=}\n") Output: x2 (computed with the non-jitted function) is correct; x3 just has zero-padding instead of a column of 99 x=Array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19]], dtype=int32) x2=Array([[ 0, 1, 99, 2, 3], [ 4, 5, 99, 6, 7], [ 8, 9, 99, 10, 11], [12, 13, 99, 14, 15], [16, 17, 99, 18, 19]], dtype=int32) x3=Array([[ 0, 1, 2, 3, 0], [ 4, 5, 6, 7, 0], [ 8, 9, 10, 11, 0], [12, 13, 14, 15, 0], [16, 17, 18, 19, 0]], dtype=int32) The same code run on a non-metal machine gives the correct results: x=Array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19]], dtype=int32) x2=Array([[ 0, 1, 99, 2, 3], [ 4, 5, 99, 6, 7], [ 8, 9, 99, 10, 11], [12, 13, 99, 14, 15], [16, 17, 99, 18, 19]], dtype=int32) x3=Array([[ 0, 1, 99, 2, 3], [ 4, 5, 99, 6, 7], [ 8, 9, 99, 10, 11], [12, 13, 99, 14, 15], [16, 17, 99, 18, 19]], dtype=int32) Not sure if this is the correct channel for bug reports, please feel free to let me know if there's a more appropriate place!

Hello iOS Developer Community, I hope this message finds you healthy and happy. I am reaching out to seek your expertise and assistance with a particular challenge I’ve encountered while using the Speak Screen and Speak Selection features on iOS. As you may know, these features are incredibly useful for reading text aloud, but they sometimes struggle with the correct pronunciation of homographs—words that are spelled the same but have different meanings and pronunciations. An example of this is the word “live,” which can be pronounced differently based on the context of the sentence. To enhance my user experience, I am looking to input corrections for the pronunciation of “live” in its “happening now” context, such as in “live broadcast” or “live event.” However, the current process requires manual entry for each phrase, which is quite labor-intensive. I am wondering if there is a way to automate or streamline this process, perhaps through a shortcut or script that allows for bulk input of these corrections. Additionally, if anyone has already compiled a list of common phrases with homographs and their correct pronunciations, I would greatly appreciate it if you could share it or guide me on where to find such resources. Your insights and guidance on this matter would be invaluable, and I believe any solutions could benefit not just myself but many other users facing similar issues. Thank you for your time and consideration. I look forward to any suggestions or advice you may have. Best regards, Alec

Will macos support amd rx7600?

Hello! We have an app that utilises the SpeechKit Framework. Especially the local on-device speech recognition for the audio files with the user selected language. Up until recently it worked as expected. However after updating one of our testing device to iOS 17.4.1 we found out that the local recognition on it stopped working completely. The error that we are getting has code 102 at its localised description reads: "Failed to access assets". That sounds just like a rear though known issue in previous iOS versions. The solution was inconvenient for our users but at least it worked – they were to go to the System settings and tweak with the dictation setting in the keyboard section. Right now no tweaks of this sort appear to help us fix the situation. We even tried to do the setting reset of the device (not the factory reset though). The error persists. it appears one one of our devices 100% of the time, halting the local recognition process. It sometimes shows on other devices for some particular languages too, but it does not show for other languages. As it is a UX breaking bug for our app, today I decided to check the logs of the Console app at the moment of the recognition attempt. There are lots of errors with code 1101 which from our research appear to be the general notifications about some local recognition setup problems. Removing the lines about the 1101 error from the log we have some interesting stuff remaining, that is (almost) never mentioned in any of the searchable webpages in the Internet. I assume they are the private API calls that the SpeechKit Framework executes under the hood: default localspeechrecognition -[UAFAssetSet assetNamed:]_block_invoke 9067C4F1-0B29-4A57-85DD-F8740DF7C344: No assets in asset set com.apple.siri.understanding default localspeechrecognition -[UAFAssetSet assetNamed:] 9067C4F1-0B29-4A57-85DD-F8740DF7C344: Returning com.apple.siri.asr.assistant from source none error localspeechrecognition -[SFEntitledAssetManager _assetWithAssetConfig:regionId:] No asset found with name: com.apple.siri.asr.assistant, asset set: com.apple.siri.understanding, usage: <private> error localspeechrecognition +[LSRConnection modelRootWithLanguage:clientID:modelOverrideURL:returningAssetType:error:] Fetch asset error (null) error localspeechrecognition -[LSRConnection prepareRecognizerWithLanguage:recognitionOverrides:modelOverrideURL:anyConfiguration:task:clientID:error:] modelRoot is nil (null) default OurApp [0x113e96d40] invalidated because the current process cancelled the connection by calling xpc_connection_cancel() Looks like there are some language-model related problems that appeared after the device was updated to 17.4.1. The Settings -> General -> Keyboard -> Dictation Languages appear to be configured correctly, the dictation toggle is On, we tried tweaking all these setting, rebooting the device and resetting the device settings. However the log lines still tell us that there is something wrong with the private resources of the SpeechKit framework. We are very concerned as the speech recognition is the core of out application's logic. And we don't understand what is the scale of possible impact of such a faulty behaviour (rare occurrences / some users / all users?) and how we can fix it to provide our users with the desired behaviour.

I was wondering if there is a quick way to convert a model trained with the open source CRFSuite for use with NLTagger? It seems like retraining should be possible but was wondering if automatic conversion was supported?

Copying from https://github.com/google/jax/issues/20750: import jax import jax.numpy as jnp def test_func(x, y): return x, y def main(): # Print available JAX devices print("JAX devices:", jax.devices()) # Create two random matrices a = jnp.array([[1.0, 2.0], [3.0, 4.0]]) b = jnp.array([[5.0, 6.0], [7.0, 8.0]]) # Perform matrix multiplication c = jnp.dot(a, b) # Print the result print("Result of matrix multiplication:") print(c) # Compute the gradient of sum of c with respect to a grad_a = jax.grad(lambda a: jnp.sum(jnp.dot(a, b)))(a) print("Gradient with respect to a:") print(grad_a) rng = jax.random.PRNGKey(0) test_input = jax.random.normal(key=rng, shape=(5,5,5)) initial_state = jax.numpy.array(0.0) x, y = jax.lax.scan(test_func, initial_state, test_input) if __name__ == "__main__": main() Gets: Platform 'METAL' is experimental and not all JAX functionality may be correctly supported! 2024-04-15 18:22:28.994752: W pjrt_plugin/src/mps_client.cc:563] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported! Metal device set to: Apple M2 Pro systemMemory: 16.00 GB maxCacheSize: 5.33 GB JAX devices: [METAL(id=0)] Result of matrix multiplication: [[19. 22.] [43. 50.]] Gradient with respect to a: [[11. 15.] [11. 15.]] zsh: segmentation fault python JAXTest.py With more info from the debugger: Current thread 0x00000001fdd3bac0 (most recent call first): File "/Users/.../anaconda3/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 1213 in __call__ My configuration is: jax-metal : 0.0.6 jax: 0.4.26 jaxlib: 0.4.23 numpy: 1.24.3 python: 3.11.8 | packaged by conda-forge | (main, Feb 16 2024, 20:49:36) [Clang 16.0.6 ] jax.devices (1 total, 1 local): [METAL(id=0)] process_count: 1 platform: uname_result(system='Darwin', root:xnu-10063.101.17~1/RELEASE_ARM64_T6020', machine='arm64') macOS 14.4.1 (23E224) Before in 3.9+0.0.3 etc it wasn't happening.

Hi, I have encountered to a segfault error when I called something via jax.lax.scan. A minimum failing example is pasted below: $ ipython Python 3.9.6 (default, Feb 3 2024, 15:58:27) Type 'copyright', 'credits' or 'license' for more information IPython 8.18.1 -- An enhanced Interactive Python. Type '?' for help. In [1]: import jax In [2]: jax.__version__ Out[2]: '0.4.22' In [3]: import jaxlib In [4]: jaxlib.__version__ Out[4]: '0.4.22' In [6]: import jax.numpy as jnp In [7]: def f(carry, x): ...: return carry + x * x, x * x ...: ...: jax.lax.scan(f, jnp.zeros((), dtype=jnp.float32), jnp.arange(3, dtype=jnp.float32)) Platform 'METAL' is experimental and not all JAX functionality may be correctly supported! 2024-04-16 01:03:52.483015: W pjrt_plugin/src/mps_client.cc:563] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported! Metal device set to: Apple M3 Max systemMemory: 36.00 GB maxCacheSize: 13.50 GB zsh: segmentation fault ipython This might be related to the thread below: https://developer.apple.com/forums/thread/749080 Strangely, when we call it jax.lax.scan is a very important building block, so I would greatly appreciate if this can be resolved soon.