Skip to content

Commit

Permalink
Remove NumPy as a hard dependency (#204)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
dgasmith authored Jul 5, 2024
1 parent eede8fa commit 94c62a0
Show file tree
Hide file tree
Showing 32 changed files with 694 additions and 628 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/Linting.yml
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ jobs:
mypy opt_einsum
black:
name: black
name: Black
runs-on: ubuntu-latest

steps:
Expand Down
11 changes: 9 additions & 2 deletions .github/workflows/Tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ on:

jobs:
miniconda-setup:
name: Env (${{ matrix.environment }}) - Py${{ matrix.python-version }}
name: Env
strategy:
fail-fast: false
matrix:
Expand All @@ -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

Expand All @@ -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: |
Expand All @@ -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
5 changes: 5 additions & 0 deletions devtools/ci_scripts/check_no_numpy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
try:
import numpy
exit(1)
except ModuleNotFoundError:
exit(0)
4 changes: 1 addition & 3 deletions devtools/conda-envs/min-deps-environment.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,10 @@ channels:
dependencies:
# Base depends
- python >=3.9
- numpy >=1.23
- nomkl

# Testing
- autoflake
- black
- black
- codecov
- isort
- mypy
Expand Down
19 changes: 19 additions & 0 deletions devtools/conda-envs/torch-only-environment.yaml
Original file line number Diff line number Diff line change
@@ -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
15 changes: 15 additions & 0 deletions opt_einsum/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
17 changes: 12 additions & 5 deletions opt_einsum/backends/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
7 changes: 3 additions & 4 deletions opt_einsum/backends/cupy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand All @@ -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
Expand Down
36 changes: 20 additions & 16 deletions opt_einsum/backends/dispatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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``,
Expand Down
6 changes: 3 additions & 3 deletions opt_einsum/backends/jax.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions opt_einsum/backends/object_arrays.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
import functools
import operator

import numpy as np

from opt_einsum.typing import ArrayType


Expand All @@ -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("->")
Expand Down
9 changes: 4 additions & 5 deletions opt_einsum/backends/tensorflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down Expand Up @@ -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)

Expand Down
7 changes: 3 additions & 4 deletions opt_einsum/backends/theano.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand All @@ -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)

Expand Down
9 changes: 4 additions & 5 deletions opt_einsum/backends/torch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 94c62a0

Please sign in to comment.