I find the same thing. The Cholesky decomposition (cho_factor) has the same problem: failed to legalize operation. Looks to me like this needs to be implemented.
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>