Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove NumPy as a hard dependency #204

Merged
merged 23 commits into from
Jul 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading