Currently (jax-metal 0.0.5) doesn't support matrix decompositions like Cholesky, LU, or Eigen, on Metal. I have a 64GB M1 Max on Sonoma 14.2.1 (23C71). Eigen raises a NotImplemented error, but Cholesky and jax.linalg.inv error out. When using the CPU by setting
export JAX_PLATFORM_NAME=CPU
then the decompositions work. The error I am getting:
XlaRuntimeError Traceback (most recent call last) Cell In[33], line 1 ----> 1 jsl.cho_factor(Sigma_0_inv)
File ~/miniconda3/envs/jaxmetal/lib/python3.10/site-packages/jax/_src/scipy/linalg.py:61, in cho_factor(failed resolving arguments) 56 @_wraps(scipy.linalg.cho_factor, 57 lax_description=_no_overwrite_and_chkfinite_doc, skip_params=('overwrite_a', 'check_finite')) 58 def cho_factor(a: ArrayLike, lower: bool = False, overwrite_a: bool = False, 59 check_finite: bool = True) -> tuple[Array, bool]: 60 del overwrite_a, check_finite # Unused ---> 61 return (cholesky(a, lower=lower), lower)
File ~/miniconda3/envs/jaxmetal/lib/python3.10/site-packages/jax/_src/scipy/linalg.py:54, in cholesky(failed resolving arguments) 49 @_wraps(scipy.linalg.cholesky, 50 lax_description=_no_overwrite_and_chkfinite_doc, skip_params=('overwrite_a', 'check_finite')) 51 def cholesky(a: ArrayLike, lower: bool = False, overwrite_a: bool = False, 52 check_finite: bool = True) -> Array: 53 del overwrite_a, check_finite # Unused ---> 54 return _cholesky(a, lower)
File ~/miniconda3/envs/jaxmetal/lib/python3.10/site-packages/jax/_src/compiler.py:255, in backend_compile(backend, module, options, host_callbacks) 250 return backend.compile(built_c, compile_options=options, 251 host_callbacks=host_callbacks) 252 # Some backends don't have host_callbacks option yet 253 # TODO(sharadmv): remove this fallback when all backends allow compile 254 # to take in host_callbacks --> 255 return backend.compile(built_c, compile_options=options)
XlaRuntimeError: UNKNOWN: /var/folders/9d/0035yr7j3bx84h3ghpp_86pc0000gn/T/ipykernel_40684/3046650730.py:1:0: error: failed to legalize operation 'mhlo.cholesky' /var/folders/9d/0035yr7j3bx84h3ghpp_86pc0000gn/T/ipykernel_40684/3046650730.py:1:0: note: see current operation: %5 = "mhlo.cholesky"(%4) {lower = true} : (tensor<200x200xf32>) -> tensor<200x200xf32>