Implement Einsum as OpFromGraph

Author: jessegrabowski

pytensor/compile/builders.py

@@ -7,7 +7,7 @@
 from typing import cast
 
 import pytensor.tensor as pt
-from pytensor import function
+from pytensor.compile.function import function
 from pytensor.compile.function.pfunc import rebuild_collect_shared
 from pytensor.compile.mode import optdb
 from pytensor.compile.sharedvalue import SharedVariable

pytensor/link/jax/dispatch/__init__.py

@@ -14,5 +14,6 @@
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.sparse
 import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.einsum
 
 # 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

@@ -4240,17 +4240,15 @@ 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
     array shape.
     """
     a = as_tensor(a)
 
-    if not isinstance(axis, tuple | list):
+    if not isinstance(axis, Sequence):
         axis = (axis,)
 
     out_ndim = len(axis) + a.ndim

pytensor/tensor/einsum.py

@@ -0,0 +1,248 @@
+import collections
+from collections.abc import Sequence
+from functools import reduce
+from itertools import pairwise
+
+from numpy.core.numeric import normalize_axis_index  # type: ignore
+
+from pytensor.compile.builders import OpFromGraph
+from pytensor.tensor.basic import (
+    arange,
+    expand_dims,
+    get_vector_length,
+    stack,
+    transpose,
+    where,
+)
+from pytensor.tensor.extra_ops import broadcast_to
+from pytensor.tensor.math import and_, eq, tensordot
+from pytensor.tensor.shape import shape_padright
+from pytensor.tensor.variable import TensorVariable
+
+
+class Einsum(OpFromGraph):
+    """
+    Wrapper Op for Einsum graphs
+    """
+
+    __props__ = ("subscripts", "optimize")
+
+    def __init__(
+        self, *args, subscripts: str, optimize: str | None = "optimal", **kwargs
+    ):
+        self.subscripts = subscripts
+        self.optimize = optimize
+        super().__init__(*args, **kwargs)
+
+
+def _iota(shape: TensorVariable, axis: int) -> TensorVariable:
+    axis = normalize_axis_index(axis, get_vector_length(shape))
+    values = arange(shape[axis])
+    return broadcast_to(shape_padright(values, axis), shape)
+
+
+def _delta(shape, axes: Sequence[int]) -> TensorVariable:
+    """This utility function exists for creating Kronecker delta arrays."""
+    base_shape = stack([shape[axis] for axis in axes])
+    iotas = [_iota(base_shape, i) for i in range(len(axes))]
+    eyes = [eq(i1, i2) for i1, i2 in pairwise(iotas)]
+    result = reduce(and_, eyes)
+    return broadcast_to(expand_dims(result, tuple(axes)), shape)
+
+
+def _removechars(s, chars):
+    return s.translate(str.maketrans(dict.fromkeys(chars)))
+
+
+def einsum(subscripts: str, *operands):
+    """
+    Multiplication and summation of tensors using the Einstein summation convention.
+
+    # TODO: Write docs
+
+    Parameters
+    ----------
+    subscripts: str
+
+    operands: sequence of TensorVariable
+        Tensors to be multiplied and summed.
+
+    Returns
+    -------
+    TensorVariable
+        The result of the einsum operation.
+    """
+    # TODO: Is this doing something clever about unknown shapes?
+    # contract_path = _poly_einsum_handlers.get(ty, _default_poly_einsum_handler)
+    # using einsum_call=True here is an internal api for opt_einsum... sorry
+
+    # TODO: Handle None static shapes
+    # TODO: Do we need this as dependency?
+    from opt_einsum import contract_path
+
+    shapes = [operand.type.shape for operand in operands]
+
+    # TODE: Do fast path at creation time, and optimize only in fast_run
+    _, contraction_list = contract_path(
+        subscripts,
+        *shapes,
+        einsum_call=True,
+        use_blas=True,
+        optimize="optimal",
+        shapes=True,
+    )
+
+    def sum_uniques(
+        operand: TensorVariable, names: str, uniques: list[str]
+    ) -> tuple[TensorVariable, str]:
+        if uniques:
+            axes = [names.index(name) for name in uniques]
+            operand = operand.sum(axes)
+            names = _removechars(names, uniques)
+        return operand, names
+
+    def sum_repeats(
+        operand: TensorVariable,
+        names: str,
+        counts: collections.Counter,
+        keep_names: str,
+    ) -> tuple[TensorVariable, str]:
+        for name, count in counts.items():
+            if count > 1:
+                axes = [i for i, n in enumerate(names) if n == name]
+                eye = _delta(operand.shape, axes)
+                operand = where(eye, operand, operand.zeros_like())
+                if name not in keep_names:
+                    operand = operand.sum(axes)
+                    names = names.replace(name, "")
+                else:
+                    operand = operand.sum(axes[:-1])
+                    names = names.replace(name, "", count - 1)
+        return operand, names
+
+    # def filter_singleton_dims(operand, names, other_shape, other_names):
+    #     eq = core.definitely_equal
+    #     keep = [
+    #         not eq(operand.shape[i], 1) or j == -1 or eq(other_shape[j], 1)
+    #         for i, j in enumerate(map(other_names.find, names))
+    #     ]
+    #     sqez_axes, keep_axes = partition_list(keep, list(range(operand.ndim)))
+    #     return lax.squeeze(operand, sqez_axes), "".join(names[i] for i in keep_axes)
+
+    einsum_operands = list(operands)  # So we can pop
+    for operand_indices, contracted_names, einstr, _, _ in contraction_list:
+        contracted_names = sorted(contracted_names)
+        assert len(contracted_names) == len(
+            set(contracted_names)
+        ), "The set was needed!"
+
+        input_str, result_names = einstr.split("->")
+        input_names = input_str.split(",")
+
+        # switch on the number of operands to be processed in this loop iteration.
+        # every case here sets 'operand' and 'names'.
+        if len(operand_indices) == 1:
+            operand = einsum_operands.pop(operand_indices[0])
+            (names,) = input_names
+            counts = collections.Counter(names)
+
+            # sum out unique contracted indices with a single reduce-sum
+            uniques = [name for name in contracted_names if counts[name] == 1]
+            operand, names = sum_uniques(operand, names, uniques)
+
+            # for every repeated index, do a contraction against an identity matrix
+            operand, names = sum_repeats(operand, names, counts, result_names)
+
+        elif len(operand_indices) == 2:
+            lhs, rhs = map(einsum_operands.pop, operand_indices)
+            lhs_names, rhs_names = input_names
+
+            # handle cases where one side of a contracting or batch dimension is 1
+            # but its counterpart is not.
+            # lhs, lhs_names = filter_singleton_dims(lhs, lhs_names, shape(rhs),
+            #                                        rhs_names)
+            # rhs, rhs_names = filter_singleton_dims(rhs, rhs_names, shape(lhs),
+            #                                        lhs_names)
+
+            lhs_counts = collections.Counter(lhs_names)
+            rhs_counts = collections.Counter(rhs_names)
+
+            # sum out unique contracted indices in lhs and rhs
+            lhs_uniques = [
+                name
+                for name in contracted_names
+                if lhs_counts[name] == 1 and rhs_counts[name] == 0
+            ]
+            lhs, lhs_names = sum_uniques(lhs, lhs_names, lhs_uniques)
+
+            rhs_uniques = [
+                name
+                for name in contracted_names
+                if rhs_counts[name] == 1 and lhs_counts[name] == 0
+            ]
+            rhs, rhs_names = sum_uniques(rhs, rhs_names, rhs_uniques)
+
+            # for every repeated index, contract against an identity matrix
+            lhs, lhs_names = sum_repeats(
+                lhs, lhs_names, lhs_counts, result_names + rhs_names
+            )
+            rhs, rhs_names = sum_repeats(
+                rhs, rhs_names, rhs_counts, result_names + lhs_names
+            )
+
+            lhs_or_rhs_names = set(lhs_names) | set(rhs_names)
+            contracted_names = [x for x in contracted_names if x in lhs_or_rhs_names]
+            lhs_and_rhs_names = set(lhs_names) & set(rhs_names)
+            batch_names = [x for x in result_names if x in lhs_and_rhs_names]
+
+            if batch_names:
+                lhs_batch, rhs_batch = tuple(
+                    zip(*[(lhs_names.find(n), rhs_names.find(n)) for n in batch_names])
+                )
+                if lhs_batch or rhs_batch:
+                    raise NotImplementedError("Batch dimensions are not yet supported")
+            else:
+                lhs_batch = rhs_batch = ()
+
+            # contract using dot_general
+            batch_names_str = "".join(batch_names)
+            if contracted_names:
+                lhs_cont, rhs_cont = tuple(
+                    zip(
+                        *[
+                            (lhs_names.index(n), rhs_names.index(n))
+                            for n in contracted_names
+                        ]
+                    )
+                )
+            else:
+                lhs_cont = rhs_cont = ()
+            deleted_names = batch_names_str + "".join(contracted_names)
+            remaining_lhs_names = _removechars(lhs_names, deleted_names)
+            remaining_rhs_names = _removechars(rhs_names, deleted_names)
+            # Try both orders of lhs and rhs, in the hope that one of them means we
+            # don't need an explicit transpose. opt_einsum likes to contract from
+            # right to left, so we expect (rhs,lhs) to have the best chance of not
+            # needing a transpose.
+            names = batch_names_str + remaining_rhs_names + remaining_lhs_names
+            if names == result_names:
+                operand = tensordot(rhs, lhs, (rhs_cont, lhs_cont))
+            else:
+                names = batch_names_str + remaining_lhs_names + remaining_rhs_names
+                operand = tensordot(lhs, rhs, axes=(lhs_cont, rhs_cont))
+
+        # the resulting 'operand' with axis labels 'names' should be a permutation of the desired result
+        assert len(names) == len(result_names) == len(set(names))
+        assert set(names) == set(result_names)
+        if names != result_names:
+            perm = tuple(names.index(name) for name in result_names)
+            operand = transpose(operand, perm)
+        einsum_operands.append(operand)  # used in next iteration
+
+    [einsum_result] = einsum_operands
+
+    return Einsum(
+        subscripts=subscripts,
+        inputs=list(operands),
+        outputs=[einsum_result],
+    )(*operands)

tests/link/jax/test_einsum.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+
+import pytensor
+import pytensor.tensor as pt
+
+
+jax = pytest.importorskip("jax")
+
+
+def test_jax_einsum():
+    subscripts = "ij, jk, kl -> il"
+    x = np.random.rand(3, 5)
+    y = np.random.rand(5, 2)
+    z = np.random.rand(2, 4)
+
+    shapes = ((3, 5), (5, 2), (2, 4))
+    x_pt, y_pt, z_pt = (
+        pt.tensor(name, shape=shape) for name, shape in zip("xyz", shapes)
+    )
+    out = pt.einsum(subscripts, x_pt, y_pt, z_pt)
+    f = pytensor.function([x_pt, y_pt, z_pt], out, mode="JAX")
+
+    np.testing.assert_allclose(f(x, y, z), np.einsum(subscripts, x, y, z))
+
+
+@pytest.mark.xfail(raises=NotImplementedError)
+def test_ellipsis_einsum():
+    subscripts = "...i,...i->..."
+    x = np.random.rand(2, 5)
+    y = np.random.rand(2, 5)
+
+    x_pt = pt.tensor("x", shape=x.shape)
+    y_pt = pt.tensor("y", shape=y.shape)
+    out = pt.einsum(subscripts, x_pt, y_pt)
+    f = pytensor.function([x_pt, y_pt], out, mode="JAX")
+
+    np.testing.assert_allclose(f(x, y), np.einsum(subscripts, x, y))