Implement Einsum as OpFromGraph

Co-authored-by: Jesse Grabowski <[email protected]>
Co-authored-by: Rob Zinkov <[email protected]>

Author: 2 people

pytensor/link/jax/dispatch/__init__.py

@@ -14,6 +14,7 @@
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.sparse
 import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.einsum
 import pytensor.link.jax.dispatch.sort
 
 # isort: on

pytensor/link/jax/dispatch/einsum.py

@@ -0,0 +1,15 @@
+import jax.numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.einsum import Einsum
+
+
+@jax_funcify.register(Einsum)
+def jax_funcify_Einsum(op, **kwargs):
+    subscripts = op.subscripts
+    optimize = op.optimize
+
+    def einsum(*operands):
+        return jnp.einsum(subscripts, *operands, optimize=optimize)
+
+    return einsum

pytensor/tensor/__init__.py

@@ -153,5 +153,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 from pytensor.tensor.functional import vectorize
 # isort: on
 
+from pytensor.tensor.einsum import einsum
+
 
 __all__ = ["random"]  # noqa: F405

pytensor/tensor/basic.py

@@ -1997,7 +1997,12 @@ def transpose(x, axes=None):
     _x = as_tensor_variable(x)
 
     if axes is None:
-        axes = list(range((_x.type.ndim - 1), -1, -1))
+        axes = tuple(range((_x.type.ndim - 1), -1, -1))
+
+    if tuple(axes) == tuple(range(len(axes))):
+        # No-op
+        return _x
+
     ret = DimShuffle(tuple(s == 1 for s in _x.type.shape), axes)(_x)
 
     if _x.name and axes == list(range((_x.type.ndim - 1), -1, -1)):
@@ -3976,6 +3981,10 @@ def moveaxis(
     source = normalize_axis_tuple(source, a.ndim, "source")
     destination = normalize_axis_tuple(destination, a.ndim, "destination")
 
+    if source == destination:
+        # It's a no-op
+        return a
+
     if len(source) != len(destination):
         raise ValueError(
             "`source` and `destination` arguments must have the same number of elements"
@@ -4290,9 +4299,7 @@ def atleast_Nd(
 atleast_3d = partial(atleast_Nd, n=3)
 
 
-def expand_dims(
-    a: np.ndarray | TensorVariable, axis: tuple[int, ...]
-) -> TensorVariable:
+def expand_dims(a: np.ndarray | TensorVariable, axis: Sequence[int]) -> TensorVariable:
     """Expand the shape of an array.
 
     Insert a new axis that will appear at the `axis` position in the expanded
@@ -4311,7 +4318,7 @@ def expand_dims(
     """
     a = as_tensor(a)
 
-    if not isinstance(axis, tuple | list):
+    if not isinstance(axis, Sequence):
         axis = (axis,)
 
     out_ndim = len(axis) + a.ndim