Implement Einsum as OpFromGraph

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

Author: 2 people

.github/workflows/test.yml

@@ -149,7 +149,7 @@ jobs:
         shell: micromamba-shell {0}
         run: |
 
-          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
+          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy opt_einsum pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 -c pytorch -c nvidia; fi
@@ -209,7 +209,7 @@ jobs:
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
-          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark  pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
+          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy opt_einsum pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark  pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'

environment.yml

@@ -11,6 +11,7 @@ dependencies:
   - compilers
   - numpy>=1.17.0,<2
   - scipy>=0.14,<1.14.0
+  - opt_einsum
   - filelock
   - etuples
   - logical-unification

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=0.14,<1.14",
     "numpy>=1.17.0,<2",
+    "opt_einsum",
     "filelock",
     "etuples",
     "logical-unification",

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