diff --git a/devtools/conda-envs/min-deps-environment.yaml b/devtools/conda-envs/min-deps-environment.yaml index e6ea382..31b733a 100644 --- a/devtools/conda-envs/min-deps-environment.yaml +++ b/devtools/conda-envs/min-deps-environment.yaml @@ -4,8 +4,6 @@ channels: dependencies: # Base depends - python >=3.9 - - numpy >=1.23 - - nomkl # Testing - autoflake diff --git a/opt_einsum/__init__.py b/opt_einsum/__init__.py index 828fc52..4e73ace 100644 --- a/opt_einsum/__init__.py +++ b/opt_einsum/__init__.py @@ -9,6 +9,8 @@ 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", "RadnomGreedy", "BranchBound", "DynamicProgramming", "shared_intermediates"] + # Handle versioneer from opt_einsum._version import get_versions # isort:skip diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py index e44d12d..43ba392 100644 --- a/opt_einsum/contract.py +++ b/opt_einsum/contract.py @@ -2,7 +2,6 @@ Contains the primary optimization and contraction routines. """ -from collections import namedtuple from decimal import Decimal from functools import lru_cache from typing import Any, Collection, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union, overload @@ -10,6 +9,7 @@ from opt_einsum import backends, blas, helpers, parser, paths, sharing from opt_einsum.typing import ( ArrayIndexType, + ArrayShaped, ArrayType, BackendType, ContractionListType, @@ -957,14 +957,11 @@ def __str__(self) -> str: return "".join(s) -Shaped = namedtuple("Shaped", ["shape"]) - - -def shape_only(shape: TensorShapeType) -> Shaped: +def shape_only(shape: TensorShapeType) -> ArrayShaped: """Dummy ``numpy.ndarray`` which has a shape only - for generating contract expressions. """ - return Shaped(shape) + return ArrayShaped(shape) # Overlaod for contract(einsum_string, *operands) diff --git a/opt_einsum/helpers.py b/opt_einsum/helpers.py index 05c82bb..5775259 100644 --- a/opt_einsum/helpers.py +++ b/opt_einsum/helpers.py @@ -2,10 +2,9 @@ Contains helper functions for opt_einsum testing scripts """ -from typing import Any, Collection, Dict, FrozenSet, Iterable, List, Literal, Optional, Tuple, Union, overload +from typing import Any, Collection, Dict, FrozenSet, Iterable, List, Tuple, overload -from opt_einsum.parser import get_symbol -from opt_einsum.typing import ArrayIndexType, ArrayType, PathType +from opt_einsum.typing import ArrayIndexType, ArrayType __all__ = ["build_views", "compute_size_by_dict", "find_contraction", "flop_count"] @@ -14,41 +13,6 @@ _default_dim_dict = {c: s for c, s in zip(_valid_chars, _sizes)} -def build_views(string: str, dimension_dict: Optional[Dict[str, int]] = None) -> List[np.ndarray]: - """ - Builds random numpy arrays for testing. - - Parameters - ---------- - string : str - List of tensor strings to build - dimension_dict : dictionary - Dictionary of index _sizes - - Returns - ------- - ret : list of np.ndarry's - The resulting views. - - Examples - -------- - >>> view = build_views('abbc', {'a': 2, 'b':3, 'c':5}) - >>> view[0].shape - (2, 3, 3, 5) - - """ - - if dimension_dict is None: - dimension_dict = _default_dim_dict - - views = [] - terms = string.split("->")[0].split(",") - for term in terms: - dims = [dimension_dict[x] for x in term] - views.append(np.random.rand(*dims)) - return views - - @overload def compute_size_by_dict(indices: Iterable[int], idx_dict: List[int]) -> int: ... @@ -194,139 +158,3 @@ def has_array_interface(array: ArrayType) -> ArrayType: return True else: return False - - -@overload -def rand_equation( - n: int, - regularity: int, - n_out: int = ..., - d_min: int = ..., - d_max: int = ..., - seed: Optional[int] = ..., - global_dim: bool = ..., - *, - return_size_dict: Literal[True], -) -> Tuple[str, PathType, Dict[str, int]]: ... - - -@overload -def rand_equation( - n: int, - regularity: int, - n_out: int = ..., - d_min: int = ..., - d_max: int = ..., - seed: Optional[int] = ..., - global_dim: bool = ..., - return_size_dict: Literal[False] = ..., -) -> Tuple[str, PathType]: ... - - -def rand_equation( - n: int, - regularity: int, - n_out: int = 0, - d_min: int = 2, - d_max: int = 9, - seed: Optional[int] = None, - global_dim: bool = False, - return_size_dict: bool = False, -) -> Union[Tuple[str, PathType, Dict[str, int]], Tuple[str, PathType]]: - """Generate a random contraction and shapes. - - Parameters: - n: Number of array arguments. - regularity: 'Regularity' of the contraction graph. This essentially determines how - many indices each tensor shares with others on average. - n_out: Number of output indices (i.e. the number of non-contracted indices). - Defaults to 0, i.e., a contraction resulting in a scalar. - d_min: Minimum dimension size. - d_max: Maximum dimension size. - seed: If not None, seed numpy's random generator with this. - global_dim: Add a global, 'broadcast', dimension to every operand. - return_size_dict: Return the mapping of indices to sizes. - - Returns: - eq: The equation string. - shapes: The array shapes. - size_dict: The dict of index sizes, only returned if ``return_size_dict=True``. - - Examples: - ```python - >>> eq, shapes = rand_equation(n=10, regularity=4, n_out=5, seed=42) - >>> eq - 'oyeqn,tmaq,skpo,vg,hxui,n,fwxmr,hitplcj,kudlgfv,rywjsb->cebda' - - >>> shapes - [(9, 5, 4, 5, 4), - (4, 4, 8, 5), - (9, 4, 6, 9), - (6, 6), - (6, 9, 7, 8), - (4,), - (9, 3, 9, 4, 9), - (6, 8, 4, 6, 8, 6, 3), - (4, 7, 8, 8, 6, 9, 6), - (9, 5, 3, 3, 9, 5)] - ``` - """ - - if seed is not None: - np.random.seed(seed) - - # total number of indices - num_inds = n * regularity // 2 + n_out - inputs = ["" for _ in range(n)] - output = [] - - size_dict = {get_symbol(i): np.random.randint(d_min, d_max + 1) for i in range(num_inds)} - - # generate a list of indices to place either once or twice - def gen(): - for i, ix in enumerate(size_dict): - # generate an outer index - if i < n_out: - output.append(ix) - yield ix - # generate a bond - else: - yield ix - yield ix - - # add the indices randomly to the inputs - for i, ix in enumerate(np.random.permutation(list(gen()))): - # make sure all inputs have at least one index - if i < n: - inputs[i] += ix - else: - # don't add any traces on same op - where = np.random.randint(0, n) - while ix in inputs[where]: - where = np.random.randint(0, n) - - inputs[where] += ix - - # possibly add the same global dim to every arg - if global_dim: - gdim = get_symbol(num_inds) - size_dict[gdim] = np.random.randint(d_min, d_max + 1) - for i in range(n): - inputs[i] += gdim - output += gdim - - # randomly transpose the output indices and form equation - output = "".join(np.random.permutation(output)) # type: ignore - eq = "{}->{}".format(",".join(inputs), output) - - # make the shapes - shapes = [tuple(size_dict[ix] for ix in op) for op in inputs] - - if return_size_dict: - return ( - eq, - shapes, - size_dict, - ) - else: - return (eq, shapes) diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py index 5505969..0354c30 100644 --- a/opt_einsum/parser.py +++ b/opt_einsum/parser.py @@ -265,32 +265,29 @@ def parse_einsum_input(operands: Any, shapes: bool = False) -> Tuple[str, str, L """ A reproduction of einsum c side einsum parsing in python. - **Parameters:** - Intakes the same inputs as `contract_path`, but NOT the keyword args. The only - supported keyword argument is: - - **shapes** - *(bool, optional)* Whether ``parse_einsum_input`` should assume arrays (the default) or - array shapes have been supplied. + Parameters: + operands: Intakes the same inputs as `contract_path`, but NOT the keyword args. The only + supported keyword argument is: + shapes: Whether ``parse_einsum_input`` should assume arrays (the default) or + array shapes have been supplied. Returns - ------- - input_strings : str - Parsed input strings - output_string : str - Parsed output string - operands : list of array_like - The operands to use in the numpy contraction - - Examples - -------- - The operand list is simplified to reduce printing: - - >>> a = np.random.rand(4, 4) - >>> b = np.random.rand(4, 4, 4) - >>> parse_einsum_input(('...a,...a->...', a, b)) - ('za,xza', 'xz', [a, b]) - - >>> parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0])) - ('za,xza', 'xz', [a, b]) + input_strings: Parsed input strings + output_string: Parsed output string + operands: The operands to use in the numpy contraction + + Examples: + The operand list is simplified to reduce printing: + + ```python + >>> a = np.random.rand(4, 4) + >>> b = np.random.rand(4, 4, 4) + >>> parse_einsum_input(('...a,...a->...', a, b)) + ('za,xza', 'xz', [a, b]) + + >>> parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0])) + ('za,xza', 'xz', [a, b]) + ``` """ if len(operands) == 0: diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py index d70f297..f598d25 100644 --- a/opt_einsum/paths.py +++ b/opt_einsum/paths.py @@ -13,8 +13,6 @@ from typing import Counter as CounterType from typing import Dict, FrozenSet, Generator, List, Optional, Sequence, Set, Tuple, Union -import numpy as np - from opt_einsum.helpers import compute_size_by_dict, flop_count from opt_einsum.typing import ArrayIndexType, PathSearchFunctionType, PathType, TensorShapeType @@ -39,17 +37,13 @@ class PathOptimizer: Subclassed optimizers should define a call method with signature: ```python - def __call__(self, inputs, output, size_dict, memory_limit=None): + def __call__(self, inputs: List[ArrayIndexType], output: ArrayIndexType, size_dict: dict[str, int], memory_limit: int | None = None) -> list[tuple[int, ...]]: \"\"\" Parameters: - inputs : list[set[str]] - The indices of each input array. - outputs : set[str] - The output indices - size_dict : dict[str, int] - The size of each index - memory_limit : int, optional - If given, the maximum allowed memory. + inputs: The indices of each input array. + outputs: The output indices + size_dict: The size of each index + memory_limit: If given, the maximum allowed memory. \"\"\" # ... compute path here ... return path diff --git a/opt_einsum/testing.py b/opt_einsum/testing.py index b3aed27..537b644 100644 --- a/opt_einsum/testing.py +++ b/opt_einsum/testing.py @@ -3,12 +3,12 @@ """ from importlib.util import find_spec -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload import pytest from opt_einsum.parser import get_symbol -from opt_einsum.typing import PathType +from opt_einsum.typing import ArrayType, PathType _valid_chars = "abcdefghijklmopqABC" _sizes = [2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3, 2, 3, 4] @@ -29,27 +29,23 @@ def import_numpy_or_skip() -> Any: return find_spec("numpy") -def build_views(string: str, dimension_dict: Optional[Dict[str, int]] = None) -> List[Any]: +def build_views(string: str, dimension_dict: Optional[Dict[str, int]] = None) -> List[ArrayType]: """ Builds random numpy arrays for testing. - Parameters - ---------- - string : str - List of tensor strings to build - dimension_dict : dictionary - Dictionary of index _sizes + Parameters: + string: List of tensor strings to build + dimension_dict: Dictionary of index _sizes Returns - ------- - ret : list of np.ndarry's The resulting views. - Examples - -------- - >>> view = build_views('abbc', {'a': 2, 'b':3, 'c':5}) - >>> view[0].shape - (2, 3, 3, 5) + Examples: + ```python + >>> view = build_views('abbc', {'a': 2, 'b':3, 'c':5}) + >>> view[0].shape + (2, 3, 3, 5) + ``` """ np = import_numpy_or_skip() @@ -65,9 +61,36 @@ def build_views(string: str, dimension_dict: Optional[Dict[str, int]] = None) -> return views +@overload def rand_equation( n: int, - reg: int, + regularity: int, + n_out: int = ..., + d_min: int = ..., + d_max: int = ..., + seed: Optional[int] = ..., + global_dim: bool = ..., + *, + return_size_dict: Literal[True], +) -> Tuple[str, PathType, Dict[str, int]]: ... + + +@overload +def rand_equation( + n: int, + regularity: int, + n_out: int = ..., + d_min: int = ..., + d_max: int = ..., + seed: Optional[int] = ..., + global_dim: bool = ..., + return_size_dict: Literal[False] = ..., +) -> Tuple[str, PathType]: ... + + +def rand_equation( + n: int, + regularity: int, n_out: int = 0, d_min: int = 2, d_max: int = 9, @@ -77,53 +100,41 @@ def rand_equation( ) -> Union[Tuple[str, PathType, Dict[str, int]], Tuple[str, PathType]]: """Generate a random contraction and shapes. - Parameters - ---------- - n : int - Number of array arguments. - reg : int - 'Regularity' of the contraction graph. This essentially determines how - many indices each tensor shares with others on average. - n_out : int, optional - Number of output indices (i.e. the number of non-contracted indices). - Defaults to 0, i.e., a contraction resulting in a scalar. - d_min : int, optional - Minimum dimension size. - d_max : int, optional - Maximum dimension size. - seed: int, optional - If not None, seed numpy's random generator with this. - global_dim : bool, optional - Add a global, 'broadcast', dimension to every operand. - return_size_dict : bool, optional - Return the mapping of indices to sizes. - - Returns - ------- - eq : str - The equation string. - shapes : list[tuple[int]] - The array shapes. - size_dict : dict[str, int] - The dict of index sizes, only returned if ``return_size_dict=True``. - - Examples - -------- - >>> eq, shapes = rand_equation(n=10, reg=4, n_out=5, seed=42) - >>> eq - 'oyeqn,tmaq,skpo,vg,hxui,n,fwxmr,hitplcj,kudlgfv,rywjsb->cebda' - - >>> shapes - [(9, 5, 4, 5, 4), - (4, 4, 8, 5), - (9, 4, 6, 9), - (6, 6), - (6, 9, 7, 8), - (4,), - (9, 3, 9, 4, 9), - (6, 8, 4, 6, 8, 6, 3), - (4, 7, 8, 8, 6, 9, 6), - (9, 5, 3, 3, 9, 5)] + Parameters: + n: Number of array arguments. + regularity: 'Regularity' of the contraction graph. This essentially determines how + many indices each tensor shares with others on average. + n_out: Number of output indices (i.e. the number of non-contracted indices). + Defaults to 0, i.e., a contraction resulting in a scalar. + d_min: Minimum dimension size. + d_max: Maximum dimension size. + seed: If not None, seed numpy's random generator with this. + global_dim: Add a global, 'broadcast', dimension to every operand. + return_size_dict: Return the mapping of indices to sizes. + + Returns: + eq: The equation string. + shapes: The array shapes. + size_dict: The dict of index sizes, only returned if ``return_size_dict=True``. + + Examples: + ```python + >>> eq, shapes = rand_equation(n=10, regularity=4, n_out=5, seed=42) + >>> eq + 'oyeqn,tmaq,skpo,vg,hxui,n,fwxmr,hitplcj,kudlgfv,rywjsb->cebda' + + >>> shapes + [(9, 5, 4, 5, 4), + (4, 4, 8, 5), + (9, 4, 6, 9), + (6, 6), + (6, 9, 7, 8), + (4,), + (9, 3, 9, 4, 9), + (6, 8, 4, 6, 8, 6, 3), + (4, 7, 8, 8, 6, 9, 6), + (9, 5, 3, 3, 9, 5)] + ``` """ np = import_numpy_or_skip() diff --git a/opt_einsum/tests/test_backends.py b/opt_einsum/tests/test_backends.py index d51d8bc..39ce5e7 100644 --- a/opt_einsum/tests/test_backends.py +++ b/opt_einsum/tests/test_backends.py @@ -4,7 +4,7 @@ import pytest from opt_einsum import backends, contract, contract_expression, sharing -from opt_einsum.contract import Shaped, infer_backend, parse_backend +from opt_einsum.contract import ArrayShaped, infer_backend, parse_backend from opt_einsum.testing import build_views try: @@ -446,7 +446,7 @@ def test_torch_with_constants(constants: Set[int]) -> None: def test_auto_backend_custom_array_no_tensordot() -> None: - x = Shaped((1, 2, 3)) + x = ArrayShaped((1, 2, 3)) # Shaped is an array-like object defined by opt_einsum - which has no TDOT assert infer_backend(x) == "opt_einsum" assert parse_backend([x], "auto") == "numpy" diff --git a/opt_einsum/tests/test_parser.py b/opt_einsum/tests/test_parser.py index 4ae6530..454ad86 100644 --- a/opt_einsum/tests/test_parser.py +++ b/opt_einsum/tests/test_parser.py @@ -18,9 +18,7 @@ def test_get_symbol() -> None: def test_parse_einsum_input() -> None: - eq = "ab,bc,cd" - ops = build_arrays_from_tuples([(2, 3), (3, 4), (4, 5)]) - input_subscripts, output_subscript, operands = parse_einsum_input([eq, *ops]) + input_subscripts, output_subscript, operands = parse_einsum_input(["ab,bc,cd", (2, 3), (3, 4), (4, 5)], shapes=True) assert input_subscripts == eq assert output_subscript == "ad" assert operands == ops @@ -39,8 +37,8 @@ def test_parse_einsum_input_shapes() -> None: import numpy as np eq = "ab,bc,cd" - shps = [(2, 3), (3, 4), (4, 5)] - input_subscripts, output_subscript, operands = parse_einsum_input([eq, *shps], shapes=True) + shapes = [(2, 3), (3, 4), (4, 5)] + input_subscripts, output_subscript, operands = parse_einsum_input([eq, *shapes], shapes=True) assert input_subscripts == eq assert output_subscript == "ad" assert np.allclose([possibly_convert_to_numpy(shp) for shp in shps], operands) diff --git a/opt_einsum/typing.py b/opt_einsum/typing.py index 175bb48..9b2d786 100644 --- a/opt_einsum/typing.py +++ b/opt_einsum/typing.py @@ -2,6 +2,7 @@ Types used in the opt_einsum package """ +from collections import namedtuple from typing import Any, Callable, Collection, Dict, FrozenSet, List, Literal, Optional, Tuple, Union TensorShapeType = Tuple[int, ...] @@ -9,6 +10,7 @@ ArrayType = Any # TODO ArrayIndexType = FrozenSet[str] +ArrayShaped = namedtuple("Shaped", ["shape"]) ContractionListType = List[Tuple[Any, ArrayIndexType, str, Optional[Tuple[str, ...]], Union[str, bool]]] PathSearchFunctionType = Callable[[List[ArrayIndexType], ArrayIndexType, Dict[str, int], Optional[int]], PathType]