Remove NumPy as a hard dependency (#204)

* Removes old tensor blas

* Reworks tests to remove numpy

* Removes numpy from helpers code

* Starts fixing up tests

* Better functional names

* Absolute paths and finishes merge conflicts

* misc

* 44 tests passing

* 44 tests passing

* 44 tests passing

* Removes numpy from core code base

* Updates MyPy

* Adds ellipses test case, closes #235, 236

* Adds a pure torch environment and potentially fixes codecov

* Adds a raw torch test

* Overhauls test_backends, allows torch to be tested without NumPy

* Removes generic shape issue

* Adds a get_shapes functionality to replace NumPy casting

* Adds edge case tests with new shape kinds

* Removes debug flag

* Attempts to fix the torch only env

* Adds a CI check to ensure there is no NumPy in the torch only env

.github/workflows/Linting.yml

@@ -47,7 +47,7 @@ jobs:
         mypy opt_einsum
 
   black:
-    name: black
+    name: Black
     runs-on: ubuntu-latest
 
     steps:

.github/workflows/Tests.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   miniconda-setup:
-    name: Env (${{ matrix.environment }}) - Py${{ matrix.python-version }}
+    name: Env
     strategy:
       fail-fast: false
       matrix:
@@ -21,6 +21,8 @@ jobs:
            environment: "min-ver"
          - python-version: 3.11
            environment: "full"
+         - python-version: 3.12
+           environment: "torch-only"
 
     runs-on: ubuntu-latest
 
@@ -43,6 +45,11 @@ jobs:
         conda config --show-sources
         conda config --show
 
+    - name: Check no NumPy for torch-only environment
+      if: matrix.environment == 'torch-only'
+      run: |
+        python devtools/ci_scripts/check_no_numpy.py
+
     - name: Install
       shell: bash -l {0}
       run: |
@@ -58,7 +65,7 @@ jobs:
       run: |
         coverage report
 
-    - uses: codecov/codecov-action@v1
+    - uses: codecov/codecov-action@v4
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
         files: ./coverage.xml

devtools/ci_scripts/check_no_numpy.py

@@ -0,0 +1,5 @@
+try:
+    import numpy
+    exit(1)
+except ModuleNotFoundError:
+    exit(0)

devtools/conda-envs/min-deps-environment.yaml

@@ -4,12 +4,10 @@ channels:
 dependencies:
     # Base depends
   - python >=3.9
-  - numpy >=1.23
-  - nomkl
 
     # Testing
   - autoflake
-  - black
+  - black 
   - codecov
   - isort
   - mypy

devtools/conda-envs/torch-only-environment.yaml

@@ -0,0 +1,19 @@
+name: test
+channels:
+  - pytorch
+  - conda-forge
+dependencies:
+    # Base depends
+  - python >=3.9
+  - pytorch::pytorch >=2.0,<3.0.0a
+  - pytorch::cpuonly
+  - mkl
+
+    # Testing
+  - autoflake
+  - black 
+  - codecov
+  - isort
+  - mypy
+  - pytest
+  - pytest-cov

opt_einsum/__init__.py

@@ -9,6 +9,21 @@
 from opt_einsum.paths import BranchBound, DynamicProgramming
 from opt_einsum.sharing import shared_intermediates
 
+__all__ = [
+    "blas",
+    "helpers",
+    "path_random",
+    "paths",
+    "contract",
+    "contract_expression",
+    "contract_path",
+    "get_symbol",
+    "RandomGreedy",
+    "BranchBound",
+    "DynamicProgramming",
+    "shared_intermediates",
+]
+
 # Handle versioneer
 from opt_einsum._version import get_versions  # isort:skip
 

opt_einsum/backends/__init__.py

@@ -3,11 +3,18 @@
 """
 
 # Backends
-from .cupy import to_cupy
-from .dispatch import build_expression, evaluate_constants, get_func, has_backend, has_einsum, has_tensordot
-from .tensorflow import to_tensorflow
-from .theano import to_theano
-from .torch import to_torch
+from opt_einsum.backends.cupy import to_cupy
+from opt_einsum.backends.dispatch import (
+    build_expression,
+    evaluate_constants,
+    get_func,
+    has_backend,
+    has_einsum,
+    has_tensordot,
+)
+from opt_einsum.backends.tensorflow import to_tensorflow
+from opt_einsum.backends.theano import to_theano
+from opt_einsum.backends.torch import to_torch
 
 __all__ = [
     "get_func",

opt_einsum/backends/cupy.py

@@ -2,9 +2,8 @@
 Required functions for optimized contractions of numpy arrays using cupy.
 """
 
-import numpy as np
-
-from ..sharing import to_backend_cache_wrap
+from opt_einsum.helpers import has_array_interface
+from opt_einsum.sharing import to_backend_cache_wrap
 
 __all__ = ["to_cupy", "build_expression", "evaluate_constants"]
 
@@ -13,7 +12,7 @@
 def to_cupy(array):  # pragma: no cover
     import cupy
 
-    if isinstance(array, np.ndarray):
+    if has_array_interface(array):
         return cupy.asarray(array)
 
     return array

opt_einsum/backends/dispatch.py

@@ -5,16 +5,14 @@
 """
 
 import importlib
-from typing import Any, Dict
+from typing import Any, Dict, Tuple
 
-import numpy
-
-from . import cupy as _cupy
-from . import jax as _jax
-from . import object_arrays
-from . import tensorflow as _tensorflow
-from . import theano as _theano
-from . import torch as _torch
+from opt_einsum.backends import cupy as _cupy
+from opt_einsum.backends import jax as _jax
+from opt_einsum.backends import object_arrays
+from opt_einsum.backends import tensorflow as _tensorflow
+from opt_einsum.backends import theano as _theano
+from opt_einsum.backends import torch as _torch
 
 __all__ = [
     "get_func",
@@ -57,16 +55,22 @@ def _import_func(func: str, backend: str, default: Any = None) -> Any:
 
 # manually cache functions as python2 doesn't support functools.lru_cache
 #     other libs will be added to this if needed, but pre-populate with numpy
-_cached_funcs = {
-    ("tensordot", "numpy"): numpy.tensordot,
-    ("transpose", "numpy"): numpy.transpose,
-    ("einsum", "numpy"): numpy.einsum,
-    # also pre-populate with the arbitrary object backend
-    ("tensordot", "object"): numpy.tensordot,
-    ("transpose", "object"): numpy.transpose,
+_cached_funcs: Dict[Tuple[str, str], Any] = {
     ("einsum", "object"): object_arrays.object_einsum,
 }
 
+try:
+    import numpy as np
+
+    _cached_funcs[("tensordot", "numpy")] = np.tensordot
+    _cached_funcs[("transpose", "numpy")] = np.transpose
+    _cached_funcs[("einsum", "numpy")] = np.einsum
+    # also pre-populate with the arbitrary object backend
+    _cached_funcs[("tensordot", "object")] = np.tensordot
+    _cached_funcs[("transpose", "object")] = np.transpose
+except ModuleNotFoundError:
+    pass
+
 
 def get_func(func: str, backend: str = "numpy", default: Any = None) -> Any:
     """Return ``{backend}.{func}``, e.g. ``numpy.einsum``,

opt_einsum/backends/jax.py

@@ -2,9 +2,7 @@
 Required functions for optimized contractions of numpy arrays using jax.
 """
 
-import numpy as np
-
-from ..sharing import to_backend_cache_wrap
+from opt_einsum.sharing import to_backend_cache_wrap
 
 __all__ = ["build_expression", "evaluate_constants"]
 
@@ -33,6 +31,8 @@ def build_expression(_, expr):  # pragma: no cover
     jax_expr = jax.jit(expr._contract)
 
     def jax_contract(*arrays):
+        import numpy as np
+
         return np.asarray(jax_expr(arrays))
 
     return jax_contract

opt_einsum/backends/object_arrays.py

@@ -5,8 +5,6 @@
 import functools
 import operator
 
-import numpy as np
-
 from opt_einsum.typing import ArrayType
 
 
@@ -31,6 +29,7 @@ def object_einsum(eq: str, *arrays: ArrayType) -> ArrayType:
     out : numpy.ndarray
         The output tensor, with ``dtype=object``.
     """
+    import numpy as np
 
     # when called by ``opt_einsum`` we will always be given a full eq
     lhs, output = eq.split("->")

opt_einsum/backends/tensorflow.py

@@ -2,9 +2,8 @@
 Required functions for optimized contractions of numpy arrays using tensorflow.
 """
 
-import numpy as np
-
-from ..sharing import to_backend_cache_wrap
+from opt_einsum.helpers import has_array_interface
+from opt_einsum.sharing import to_backend_cache_wrap
 
 __all__ = ["to_tensorflow", "build_expression", "evaluate_constants"]
 
@@ -40,13 +39,13 @@ def to_tensorflow(array, constant=False):
     tf, device, eager = _get_tensorflow_and_device()
 
     if eager:
-        if isinstance(array, np.ndarray):
+        if has_array_interface(array):
             with tf.device(device):
                 return tf.convert_to_tensor(array)
 
         return array
 
-    if isinstance(array, np.ndarray):
+    if has_array_interface(array):
         if constant:
             return tf.convert_to_tensor(array)
 

opt_einsum/backends/theano.py

@@ -2,9 +2,8 @@
 Required functions for optimized contractions of numpy arrays using theano.
 """
 
-import numpy as np
-
-from ..sharing import to_backend_cache_wrap
+from opt_einsum.helpers import has_array_interface
+from opt_einsum.sharing import to_backend_cache_wrap
 
 __all__ = ["to_theano", "build_expression", "evaluate_constants"]
 
@@ -14,7 +13,7 @@ def to_theano(array, constant=False):
     """Convert a numpy array to ``theano.tensor.TensorType`` instance."""
     import theano
 
-    if isinstance(array, np.ndarray):
+    if has_array_interface(array):
         if constant:
             return theano.tensor.constant(array)
 

opt_einsum/backends/torch.py

@@ -2,10 +2,9 @@
 Required functions for optimized contractions of numpy arrays using pytorch.
 """
 
-import numpy as np
-
-from ..parser import convert_to_valid_einsum_chars
-from ..sharing import to_backend_cache_wrap
+from opt_einsum.helpers import has_array_interface
+from opt_einsum.parser import convert_to_valid_einsum_chars
+from opt_einsum.sharing import to_backend_cache_wrap
 
 __all__ = [
     "transpose",
@@ -104,7 +103,7 @@ def tensordot(x, y, axes=2):
 def to_torch(array):
     torch, device = _get_torch_and_device()
 
-    if isinstance(array, np.ndarray):
+    if has_array_interface(array):
         return torch.from_numpy(array).to(device)
 
     return array