support jax 'linalg.expm'

jcmgray · Sep 29, 2020 · ecceb80 · ecceb80
1 parent 6faf4dc
commit ecceb80
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/autoray/autoray.py b/autoray/autoray.py
@@ -555,6 +555,7 @@ def jax_to_numpy(x):
 _FUNCS['jax', 'random.normal'] = jax_random_normal
 _MODULE_ALIASES['jax'] = 'jax.numpy'
 _SUBMODULE_ALIASES['jax', 'complex'] = 'jax.lax'
+_SUBMODULE_ALIASES['jax', 'linalg.expm'] = 'jax.scipy.linalg'
 _CUSTOM_WRAPPERS['jax', 'linalg.qr'] = qr_allow_fat
 _CUSTOM_WRAPPERS['jax', 'linalg.svd'] = svd_not_full_matrices_wrapper