Fix ellipsis and drop dependency on opt_einsum

Author: ricardoV94

.github/workflows/test.yml

@@ -154,7 +154,7 @@ jobs:
         shell: micromamba-shell {0}
         run: |
 
-          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
+          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock
           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
@@ -215,7 +215,7 @@ jobs:
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
-          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
+          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
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'

environment.yml

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

pyproject.toml

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

pytensor/tensor/einsum.py

@@ -6,12 +6,11 @@
 from typing import cast
 
 import numpy as np
+from numpy.core.einsumfunc import _find_contraction, _parse_einsum_input  # type: ignore
 from numpy.core.numeric import (  # type: ignore
     normalize_axis_index,
     normalize_axis_tuple,
 )
-from opt_einsum.helpers import find_contraction
-from opt_einsum.parser import parse_einsum_input
 
 from pytensor.compile.builders import OpFromGraph
 from pytensor.tensor import TensorLike
@@ -129,9 +128,6 @@ def _general_dot(
     core_lhs_axes = tuple(np.array(lhs_axes) - lhs_n_batch_axes)
     core_rhs_axes = tuple(np.array(rhs_axes) - rhs_n_batch_axes)
 
-    # TODO: tensordot produces very complicated graphs unnecessarily
-    #  In some cases we are just doing elemwise addition after some transpositions
-    #  We also have some Blockwise(Reshape) that will slow down things!
     if signature == "(),()->()":
         # Just a multiplication
         out = lhs * rhs
@@ -146,7 +142,7 @@ def _general_dot(
 PATH = tuple[tuple[int] | tuple[int, int]]
 
 
-def contraction_list_from_path(
+def _contraction_list_from_path(
     subscripts: str, operands: Sequence[TensorLike], path: PATH
 ):
     """
@@ -189,7 +185,7 @@ def contraction_list_from_path(
     fake_operands = [
         np.zeros([1 if dim == 1 else 0 for dim in x.type.shape]) for x in operands
     ]
-    input_subscripts, output_subscript, operands = parse_einsum_input(
+    input_subscripts, output_subscript, operands = _parse_einsum_input(
         (subscripts, *fake_operands)
     )
 
@@ -204,7 +200,7 @@ def contraction_list_from_path(
         # Make sure we remove inds from right to left
         contract_inds = tuple(sorted(contract_inds, reverse=True))
 
-        contract_tuple = find_contraction(contract_inds, input_sets, output_set)
+        contract_tuple = _find_contraction(contract_inds, input_sets, output_set)
         out_inds, input_sets, idx_removed, idx_contract = contract_tuple
 
         tmp_inputs = [input_list.pop(x) for x in contract_inds]
@@ -354,12 +350,6 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
 
     # 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
-
     operands = [as_tensor(operand) for operand in operands]
     shapes = [operand.type.shape for operand in operands]
 
@@ -375,20 +365,21 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
             path = [(0,)]
         else:
             path = [(1, 0) for i in range(len(operands) - 1)]
-        contraction_list = contraction_list_from_path(subscripts, operands, path)
+        contraction_list = _contraction_list_from_path(subscripts, operands, path)
 
         # If there are only 1 or 2 operands, there is no optimization to be done?
         optimize = len(operands) <= 2
     else:
         # Case 2: All operands have known shapes. In this case, we can use opt_einsum to compute the optimal
         # contraction order.
-        _, contraction_list = contract_path(
+        # Call _implementation to bypass dispatch
+        _, contraction_list = np.einsum_path(
             subscripts,
-            *shapes,
+            # Numpy einsum_path requires arrays even though only the shapes matter
+            # It's not trivial to duck-type our way around because of internal call to `asanyarray`
+            *[np.empty(shape) for shape in shapes],
             einsum_call=True,
-            use_blas=True,
             optimize="optimal",
-            shapes=True,
         )
         path = [contraction[0] for contraction in contraction_list]
         optimize = True