From c15aec2c7dbbcaf5b9790b980f699bf295baa7d9 Mon Sep 17 00:00:00 2001 From: Daniel Smith Date: Thu, 26 Sep 2024 09:37:32 -0400 Subject: [PATCH] v3.4 Release notes and removes NumPy kwargs (#239) * Removes explicit kwargs * Adds 3.4 changelog * Pins MyPy * Type fix * Fixes typo in * Allows booleans as optimize, closes #219 * Update docs/changelog.md Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> * Update docs/changelog.md Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> --------- Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> --- devtools/conda-envs/full-environment.yaml | 2 +- devtools/conda-envs/min-deps-environment.yaml | 4 +- devtools/conda-envs/min-ver-environment.yaml | 2 +- .../conda-envs/torch-only-environment.yaml | 2 +- docs/changelog.md | 29 ++++++++ opt_einsum/contract.py | 66 ++++++------------- opt_einsum/paths.py | 9 ++- opt_einsum/tests/test_contract.py | 13 ++++ opt_einsum/tests/test_input.py | 9 +-- opt_einsum/tests/test_paths.py | 2 +- 10 files changed, 75 insertions(+), 63 deletions(-) diff --git a/devtools/conda-envs/full-environment.yaml b/devtools/conda-envs/full-environment.yaml index 31969e77..6731ad6f 100644 --- a/devtools/conda-envs/full-environment.yaml +++ b/devtools/conda-envs/full-environment.yaml @@ -16,7 +16,7 @@ dependencies: # Testing - codecov - - mypy + - mypy ==1.11* - pytest - pytest-cov - ruff ==0.5.* diff --git a/devtools/conda-envs/min-deps-environment.yaml b/devtools/conda-envs/min-deps-environment.yaml index 4d93a531..7527073e 100644 --- a/devtools/conda-envs/min-deps-environment.yaml +++ b/devtools/conda-envs/min-deps-environment.yaml @@ -7,7 +7,7 @@ dependencies: # Testing - codecov - - mypy + - mypy ==1.11* - pytest - pytest-cov - - ruff ==0.5.* + - ruff ==0.6.* diff --git a/devtools/conda-envs/min-ver-environment.yaml b/devtools/conda-envs/min-ver-environment.yaml index c7e02ea0..3a8dd27b 100644 --- a/devtools/conda-envs/min-ver-environment.yaml +++ b/devtools/conda-envs/min-ver-environment.yaml @@ -13,7 +13,7 @@ dependencies: # Testing - codecov - - mypy + - mypy ==1.11* - pytest - pytest-cov - ruff ==0.5.* diff --git a/devtools/conda-envs/torch-only-environment.yaml b/devtools/conda-envs/torch-only-environment.yaml index 6f622221..f65528b9 100644 --- a/devtools/conda-envs/torch-only-environment.yaml +++ b/devtools/conda-envs/torch-only-environment.yaml @@ -11,7 +11,7 @@ dependencies: # Testing - codecov - - mypy + - mypy ==1.11* - pytest - pytest-cov - ruff ==0.5.* diff --git a/docs/changelog.md b/docs/changelog.md index 16720cb7..78879694 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -1,6 +1,35 @@ Changelog ========= +## 3.4.0 / 2024-09-XX + +NumPy has been removed from `opt_einsum` as a dependency allowing for more flexible installs. + +**New Features** + + - [\#160](https://github.com/dgasmith/opt_einsum/pull/160) Migrates docs to MkDocs Material and GitHub pages hosting. + - [\#161](https://github.com/dgasmith/opt_einsum/pull/161) Adds Python type annotations to the code base. + - [\#204](https://github.com/dgasmith/opt_einsum/pull/204) Removes NumPy as a hard dependency. + +**Enhancements** + +- [\#154](https://github.com/dgasmith/opt_einsum/pull/154) Prevents an infinite recursion error when the `memory_limit` was set very low for the `dp` algorithm. +- [\#155](https://github.com/dgasmith/opt_einsum/pull/155) Adds flake8 spell check to the doc strings +- [\#159](https://github.com/dgasmith/opt_einsum/pull/159) Migrates to GitHub actions for CI. +- [\#174](https://github.com/dgasmith/opt_einsum/pull/174) Prevents double contracts of floats in dynamic paths. +- [\#196](https://github.com/dgasmith/opt_einsum/pull/196) Allows `backend=None` which is equivalent to `backend='auto'` +- [\#208](https://github.com/dgasmith/opt_einsum/pull/208) Switches to `ConfigParser` insetad of `SafeConfigParser` for Python 3.12 compatability. +- [\#228](https://github.com/dgasmith/opt_einsum/pull/228) `backend='jaxlib'` is now an alias for the `jax` library +- [\#237](https://github.com/dgasmith/opt_einsum/pull/237) Switches to `ruff` for formatting and linting. +- [\#238](https://github.com/dgasmith/opt_einsum/pull/238) Removes `numpy`-specific keyword args from being explicitly defined in `contract` and uses `**kwargs` instead. + +**Bug Fixes** + +- [\#195](https://github.com/dgasmith/opt_einsum/pull/195) Fixes a bug where `dp` would not work for scalar-only contractions. +- [\#200](https://github.com/dgasmith/opt_einsum/pull/200) Fixes a bug where `parse_einsum_input` would not correctly respect shape-only contractions. +- [\#222](https://github.com/dgasmith/opt_einsum/pull/222) Fixes an erorr in `parse_einsum_input` where an output subscript specified multiple times was not correctly caught. +- [\#229](https://github.com/dgasmith/opt_einsum/pull/229) Fixes a bug where empty contraction lists in `PathInfo` would cause an error. + ## 3.3.0 / 2020-07-19 Adds a `object` backend for optimized contractions on arbitrary Python objects. diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py index 7a7dac94..153d060d 100644 --- a/opt_einsum/contract.py +++ b/opt_einsum/contract.py @@ -26,9 +26,6 @@ ## Common types -_OrderKACF = Literal[None, "K", "A", "C", "F"] - -_Casting = Literal["no", "equiv", "safe", "same_kind", "unsafe"] _MemoryLimit = Union[None, int, Decimal, Literal["max_input"]] @@ -284,7 +281,7 @@ def contract_path( #> 5 defg,hd->efgh efgh->efgh ``` """ - if optimize is True: + if (optimize is True) or (optimize is None): optimize = "auto" # Hidden option, only einsum should call this @@ -344,9 +341,11 @@ def contract_path( naive_cost = helpers.flop_count(indices, inner_product, num_ops, size_dict) # Compute the path - if not isinstance(optimize, (str, paths.PathOptimizer)): + if optimize is False: + path_tuple: PathType = [tuple(range(num_ops))] + elif not isinstance(optimize, (str, paths.PathOptimizer)): # Custom path supplied - path_tuple: PathType = optimize # type: ignore + path_tuple = optimize # type: ignore elif num_ops <= 2: # Nothing to be optimized path_tuple = [tuple(range(num_ops))] @@ -479,9 +478,6 @@ def contract( subscripts: str, *operands: ArrayType, out: ArrayType = ..., - dtype: Any = ..., - order: _OrderKACF = ..., - casting: _Casting = ..., use_blas: bool = ..., optimize: OptimizeKind = ..., memory_limit: _MemoryLimit = ..., @@ -495,9 +491,6 @@ def contract( subscripts: ArrayType, *operands: Union[ArrayType, Collection[int]], out: ArrayType = ..., - dtype: Any = ..., - order: _OrderKACF = ..., - casting: _Casting = ..., use_blas: bool = ..., optimize: OptimizeKind = ..., memory_limit: _MemoryLimit = ..., @@ -510,9 +503,6 @@ def contract( subscripts: Union[str, ArrayType], *operands: Union[ArrayType, Collection[int]], out: Optional[ArrayType] = None, - dtype: Optional[str] = None, - order: _OrderKACF = "K", - casting: _Casting = "safe", use_blas: bool = True, optimize: OptimizeKind = True, memory_limit: _MemoryLimit = None, @@ -527,9 +517,6 @@ def contract( subscripts: Specifies the subscripts for summation. *operands: These are the arrays for the operation. out: A output array in which set the resulting output. - dtype: The dtype of the given contraction, see np.einsum. - order: The order of the resulting contraction, see np.einsum. - casting: The casting procedure for operations of different dtype, see np.einsum. use_blas: Do you use BLAS for valid operations, may use extra memory for more intermediates. optimize:- Choose the type of path the contraction will be optimized with - if a list is given uses this as the path. @@ -551,11 +538,12 @@ def contract( - `'branch-2'` An even more restricted version of 'branch-all' that only searches the best two options at each step. Scales exponentially with the number of terms in the contraction. - - `'auto'` Choose the best of the above algorithms whilst aiming to + - `'auto', None, True` Choose the best of the above algorithms whilst aiming to keep the path finding time below 1ms. - `'auto-hq'` Aim for a high quality contraction, choosing the best of the above algorithms whilst aiming to keep the path finding time below 1sec. + - `False` will not optimize the contraction. memory_limit:- Give the upper bound of the largest intermediate tensor contract will build. - None or -1 means there is no limit. @@ -586,21 +574,18 @@ def contract( performed optimally. When NumPy is linked to a threaded BLAS, potential speedups are on the order of 20-100 for a six core machine. """ - if optimize is True: + if (optimize is True) or (optimize is None): optimize = "auto" operands_list = [subscripts] + list(operands) - einsum_kwargs = {"out": out, "dtype": dtype, "order": order, "casting": casting} # If no optimization, run pure einsum if optimize is False: - return _einsum(*operands_list, **einsum_kwargs) + return _einsum(*operands_list, out=out, **kwargs) # Grab non-einsum kwargs gen_expression = kwargs.pop("_gen_expression", False) constants_dict = kwargs.pop("_constants_dict", {}) - if len(kwargs): - raise TypeError(f"Did not understand the following kwargs: {kwargs.keys()}") if gen_expression: full_str = operands_list[0] @@ -613,11 +598,9 @@ def contract( # check if performing contraction or just building expression if gen_expression: - return ContractExpression(full_str, contraction_list, constants_dict, dtype=dtype, order=order, casting=casting) + return ContractExpression(full_str, contraction_list, constants_dict, **kwargs) - return _core_contract( - operands, contraction_list, backend=backend, out=out, dtype=dtype, order=order, casting=casting - ) + return _core_contract(operands, contraction_list, backend=backend, out=out, **kwargs) @lru_cache(None) @@ -651,9 +634,7 @@ def _core_contract( backend: Optional[str] = "auto", evaluate_constants: bool = False, out: Optional[ArrayType] = None, - dtype: Optional[str] = None, - order: _OrderKACF = "K", - casting: _Casting = "safe", + **kwargs: Any, ) -> ArrayType: """Inner loop used to perform an actual contraction given the output from a ``contract_path(..., einsum_call=True)`` call. @@ -703,7 +684,7 @@ def _core_contract( axes = ((), ()) # Contract! - new_view = _tensordot(*tmp_operands, axes=axes, backend=backend) + new_view = _tensordot(*tmp_operands, axes=axes, backend=backend, **kwargs) # Build a new view if needed if (tensor_result != results_index) or handle_out: @@ -718,9 +699,7 @@ def _core_contract( out_kwarg: Union[None, ArrayType] = None if handle_out: out_kwarg = out - new_view = _einsum( - einsum_str, *tmp_operands, backend=backend, dtype=dtype, order=order, casting=casting, out=out_kwarg - ) + new_view = _einsum(einsum_str, *tmp_operands, backend=backend, out=out_kwarg, **kwargs) # Append new items and dereference what we can operands.append(new_view) @@ -768,15 +747,11 @@ def __init__( contraction: str, contraction_list: ContractionListType, constants_dict: Dict[int, ArrayType], - dtype: Optional[str] = None, - order: _OrderKACF = "K", - casting: _Casting = "safe", + **kwargs: Any, ): - self.contraction_list = contraction_list - self.dtype = dtype - self.order = order - self.casting = casting self.contraction = format_const_einsum_str(contraction, constants_dict.keys()) + self.contraction_list = contraction_list + self.kwargs = kwargs # need to know _full_num_args to parse constants with, and num_args to call with self._full_num_args = contraction.count(",") + 1 @@ -844,9 +819,7 @@ def _contract( out=out, backend=backend, evaluate_constants=evaluate_constants, - dtype=self.dtype, - order=self.order, - casting=self.casting, + **self.kwargs, ) def _contract_with_conversion( @@ -943,8 +916,7 @@ def __str__(self) -> str: for i, c in enumerate(self.contraction_list): s.append(f"\n {i + 1}. ") s.append(f"'{c[2]}'" + (f" [{c[-1]}]" if c[-1] else "")) - kwargs = {"dtype": self.dtype, "order": self.order, "casting": self.casting} - s.append(f"\neinsum_kwargs={kwargs}") + s.append(f"\neinsum_kwargs={self.kwargs}") return "".join(s) diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py index 2f469030..5c38c45d 100644 --- a/opt_einsum/paths.py +++ b/opt_einsum/paths.py @@ -499,9 +499,14 @@ def branch( output: ArrayIndexType, size_dict: Dict[str, int], memory_limit: Optional[int] = None, - **optimizer_kwargs: Dict[str, Any], + nbranch: Optional[int] = None, + cutoff_flops_factor: int = 4, + minimize: str = "flops", + cost_fn: str = "memory-removed", ) -> PathType: - optimizer = BranchBound(**optimizer_kwargs) # type: ignore + optimizer = BranchBound( + nbranch=nbranch, cutoff_flops_factor=cutoff_flops_factor, minimize=minimize, cost_fn=cost_fn + ) return optimizer(inputs, output, size_dict, memory_limit) diff --git a/opt_einsum/tests/test_contract.py b/opt_einsum/tests/test_contract.py index 313198d9..72d90f7e 100644 --- a/opt_einsum/tests/test_contract.py +++ b/opt_einsum/tests/test_contract.py @@ -14,6 +14,7 @@ # NumPy is required for the majority of this file np = pytest.importorskip("numpy") + tests = [ # Test scalar-like operations "a,->a", @@ -99,6 +100,18 @@ ] +@pytest.mark.parametrize("optimize", (True, False, None)) +def test_contract_plain_types(optimize: OptimizeKind) -> None: + expr = "ij,jk,kl->il" + ops = [np.random.rand(2, 2), np.random.rand(2, 2), np.random.rand(2, 2)] + + path = contract_path(expr, *ops, optimize=optimize) + assert len(path) == 2 + + result = contract(expr, *ops, optimize=optimize) + assert result.shape == (2, 2) + + @pytest.mark.parametrize("string", tests) @pytest.mark.parametrize("optimize", _PATH_OPTIONS) def test_compare(optimize: OptimizeKind, string: str) -> None: diff --git a/opt_einsum/tests/test_input.py b/opt_einsum/tests/test_input.py index 8ce9b0a6..24b1ead3 100644 --- a/opt_einsum/tests/test_input.py +++ b/opt_einsum/tests/test_input.py @@ -163,16 +163,9 @@ def test_value_errors(contract_fn: Any) -> None: # broadcasting to new dimensions must be enabled explicitly with pytest.raises(ValueError): contract_fn("i", np.arange(6).reshape(2, 3)) - if contract_fn is contract: - # contract_path does not have an `out` parameter - with pytest.raises(ValueError): - contract_fn("i->i", [[0, 1], [0, 1]], out=np.arange(4).reshape(2, 2)) with pytest.raises(TypeError): - contract_fn("i->i", [[0, 1], [0, 1]], bad_kwarg=True) - - with pytest.raises(ValueError): - contract_fn("i->i", [[0, 1], [0, 1]], memory_limit=-1) + contract_fn("ij->ij", [[0, 1], [0, 1]], bad_kwarg=True) @pytest.mark.parametrize( diff --git a/opt_einsum/tests/test_paths.py b/opt_einsum/tests/test_paths.py index fff2ad32..22ed1429 100644 --- a/opt_einsum/tests/test_paths.py +++ b/opt_einsum/tests/test_paths.py @@ -124,7 +124,7 @@ def test_flop_cost() -> None: def test_bad_path_option() -> None: - with pytest.raises(TypeError): + with pytest.raises(KeyError): oe.contract("a,b,c", [1], [2], [3], optimize="optimall", shapes=True) # type: ignore