add jaxlib alias

dgasmith · May 13, 2024 · 82eeae8 · 82eeae8
1 parent 57456fb
commit 82eeae8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/opt_einsum/backends/dispatch.py b/opt_einsum/backends/dispatch.py
@@ -31,6 +31,7 @@
     "theano": "theano.tensor",
     "torch": "opt_einsum.backends.torch",
     "jax": "jax.numpy",
+    "jaxlib": "jax.numpy",
     "autograd": "autograd.numpy",
     "mars": "mars.tensor",
 }