Merge pull request #26472 from jakevdp:jnp-einsum

PiperOrigin-RevId: 726580373
jax-ml · Feb 13, 2025 · 5ebb7eb · 5ebb7eb
2 parents 229aa65 + 7ab7b21
commit 5ebb7eb
Show file tree

Hide file tree

Showing 4 changed files with 586 additions and 551 deletions.
diff --git a/jax/_src/export/shape_poly.py b/jax/_src/export/shape_poly.py
@@ -42,7 +42,7 @@
 from jax._src import effects
 from jax._src.lax import lax
 from jax._src.interpreters import mlir
-from jax._src.numpy import lax_numpy
+from jax._src.numpy import einsum as jnp_einsum
 from jax._src import source_info_util
 from jax._src import tree_util
 from jax._src import util
@@ -1267,7 +1267,7 @@ def fake_dim(d):
     contract_operands.append(operands[idx[0]])
   return contract_operands, contractions
 
-lax_numpy._poly_einsum_handlers[_DimExpr] = _einsum_contract_path
+jnp_einsum._poly_einsum_handlers[_DimExpr] = _einsum_contract_path
 
 # To implement shape-constraint checking we use a shape assertion primitive.
 #    shape_assertion_p.bind(assert_what: bool, *error_message_inputs,