Merge pull request #18 from HaoTy/main

Alias jax.scipy
jcmgray · Jan 29, 2024 · f9b0a44 · f9b0a44
2 parents 364da7c + 9fab817
commit f9b0a44
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/autoray/autoray.py b/autoray/autoray.py
@@ -1495,6 +1495,7 @@ def jax_to_numpy(x):
 
 
 _BACKEND_ALIASES["jaxlib"] = "jax"
+_MODULE_ALIASES["jax.scipy"] = "jax.scipy"
 _MODULE_ALIASES["jax"] = "jax.numpy"
 _SUBMODULE_ALIASES["jax", "complex"] = "jax.lax"
 _SUBMODULE_ALIASES["jax", "linalg.expm"] = "jax.scipy.linalg"