batch-tensorsolve 0.1.0__py3-none-any.whl

batch_tensorsolve/__init__.py

@@ -0,0 +1,4 @@
+__version__ = "0.1.0"
+from ._main import AmbiguousBatchAxesWarning, broadcast_without_repeating, btensorsolve
+
+__all__ = ["AmbiguousBatchAxesWarning", "broadcast_without_repeating", "btensorsolve"]

batch_tensorsolve/_main.py

@@ -0,0 +1,178 @@
+import warnings
+from typing import TypeVar
+
+from array_api_compat import array_namespace
+
+TArray = TypeVar("TArray")
+
+
+class AmbiguousBatchAxesWarning(RuntimeWarning):
+    pass
+
+
+def broadcast_without_repeating(
+    *arrays: TArray, check_same_ndim: bool = False
+) -> tuple[TArray, ...]:
+    """
+    Broadcast arrays without repeating the data.
+
+    Parameters
+    ----------
+    arrays : TArray
+        The arrays to broadcast.
+    check_same_ndim : bool, optional
+        Whether to check if all arrays have the same number of dimensions.
+        Default is False.
+
+    Returns
+    -------
+    tuple[TArray]
+        The broadcasted arrays.
+
+    """
+    xp = array_namespace(*arrays)
+    arrays_ = tuple(xp.asarray(a) for a in arrays)
+    xp.broadcast_shapes(*[a.shape for a in arrays_])
+    if check_same_ndim:
+        if len({a.ndim for a in arrays_}) != 1:
+            raise ValueError(
+                "All arrays must have the same number of dimensions, "
+                f"but got {tuple(a.ndim for a in arrays_)}"
+            )
+        return arrays_
+    max_dim = max(a.ndim for a in arrays_)
+    return tuple(array[(None,) * (max_dim - array.ndim) + (...,)] for array in arrays_)
+
+
+def btensorsolve(
+    a: TArray,
+    b: TArray,
+    /,
+    *,
+    axes: tuple[int] | None = None,
+    num_batch_axes: int | None = None,
+) -> TArray:
+    """
+    Solve the tensor equation ``a x = b`` for x.
+
+    It is assumed that all indices of `x` are summed over in the product,
+    together with the rightmost indices of `a`, as is done in, for example,
+    ``tensordot(a, x, axes=x.ndim)``.
+
+    Parameters
+    ----------
+    a : array_like
+        Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
+        the shape of that sub-tensor of `a` consisting of the appropriate
+        number of its rightmost indices, and must be such that
+        **there exists i that**
+        ``prod(Q) == prod(b.shape[i:])``
+    b : array_like
+        Right-hand tensor, which can be of any shape.
+    axes : tuple of ints, optional
+        Axes in `a` to reorder to the right, before inversion.
+        If None (default), no reordering is done.
+    num_batch_axes : int, optional
+        The number of batch dimensions. If None (default), the number of batch
+        dimensions is inferred from the shapes of `a` and `b`.
+
+        Let ``shape = np.broadcast_shapes(a.shape[:b.ndim],
+        b.shape) + a.shape[b.ndim:]``.
+
+        It is recommended to specify this argument, as the inference
+        might be wrong if there exists i >= 0, j > 0 that
+        ``prod(shape[:i]) == prod(shape[i+j:])``. (j + 1 possibilities)
+
+        For example, if `a` has shape (3, 1, 1, 2, 2) and `b` has shape
+        (3, 1, 1, 2), it is possible that
+        - axis 1 is the batch axes and desired output shape is (3, 2)
+        - axis 1, 2 are the batch axes and desired output shape is (3, 1, 2)
+        - axis 1, 2, 3 are the batch axes and desired output shape is (3, 1, 1, 2)
+
+    Returns
+    -------
+    x : ndarray, shape Q
+
+    Raises
+    ------
+    LinAlgError
+        If `a` is singular or not 'square' (in the above sense).
+
+    Warnings
+    --------
+    AmbiguousBatchAxesWarning
+        If the number of batch axes cannot be inferred from the shapes of
+        `a` and `b`, and `num_batch_axes` is not specified.
+
+    See Also
+    --------
+    numpy.tensordot, tensorinv, numpy.einsum
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> rng = np.random.default_rng()
+    >>> a = rng.normal(size=(2, 2*3, 4, 2, 3, 4))
+    >>> b = rng.normal(size=(2, 2*3, 4))
+    >>> x = np.linalg.tensorsolve(a, b)
+    >>> x.shape
+    (2, 2, 3, 4)
+    >>> np.allclose(np.einsum('...ijklm,...klm->...ij', a, x), b)
+    True
+
+    """
+    # https://github.com/numpy/numpy/blob/
+    # e7a123b2d3eca9897843791dd698c1803d9a39c2/numpy/linalg/_linalg.py#L291
+    xp = array_namespace(a, b)
+    a_ = xp.asarray(a)
+    b_ = xp.asarray(b)
+    an = a_.ndim
+    if axes is not None:
+        allaxes = list(range(0, an))
+        for k in axes:
+            allaxes.remove(k)
+            allaxes.insert(an, k)
+        a_ = a_.transpose(allaxes)
+
+    # find right dimensions
+    # a = [2 (dim1) 2 2 3 (dim2) 2 6]
+    # b = [2 (dim1) 2 2 3 (dim2)]
+    axis_sol_last = b_.ndim
+    ashape = a_.shape
+    bshape = b_.shape
+    shape = xp.broadcast_shapes(ashape[:axis_sol_last], bshape)
+
+    # the dimension of the linear system
+    sol_size = int(xp.prod(ashape[axis_sol_last:]))
+    if num_batch_axes is None:
+        sol_size_current = 1
+        for num_batch_axes in range(axis_sol_last - 1, -1, -1):
+            sol_size_current *= shape[num_batch_axes]
+            if sol_size_current == sol_size:
+                break
+        else:
+            raise ValueError(
+                "Unable to divide batch dimensions and solution dimensions"
+            )
+
+        if num_batch_axes > 0 and shape[num_batch_axes - 1] == 1:
+            warnings.warn(
+                "It is impossible to infer the number of "
+                "batch axes from the shapes of `a` and `b`. "
+                "Consider specifying `num_batch_axes` explicitly.",
+                AmbiguousBatchAxesWarning,
+                stacklevel=2,
+            )
+
+    a_ = xp.broadcast_to(
+        a_,
+        ashape[:num_batch_axes]
+        + shape[num_batch_axes:axis_sol_last]
+        + ashape[axis_sol_last:],
+    )
+    b_ = xp.broadcast_to(b_, bshape[:num_batch_axes] + shape[num_batch_axes:])
+    a_ = a_.reshape(ashape[:num_batch_axes] + (sol_size, sol_size))
+    b_ = b_.reshape(bshape[:num_batch_axes] + (sol_size, 1))
+    x = xp.linalg.solve(a_, b_)
+    x = x.reshape(shape[:num_batch_axes] + ashape[axis_sol_last:])
+    return x

batch_tensorsolve-0.1.0.dist-info/METADATA

@@ -0,0 +1,142 @@
+Metadata-Version: 2.4
+Name: batch-tensorsolve
+Version: 0.1.0
+Summary: Batched tensorsolve() for NumPy / PyTorch / JAX
+Author-email: 34j <34j.95a2p@simplelogin.com>
+License: MIT
+Project-URL: Bug Tracker, https://github.com/34j/batch-tensorsolve/issues
+Project-URL: Changelog, https://github.com/34j/batch-tensorsolve/blob/main/CHANGELOG.md
+Project-URL: documentation, https://batch-tensorsolve.readthedocs.io
+Project-URL: repository, https://github.com/34j/batch-tensorsolve
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: array-api-compat>=1.11.2
+Requires-Dist: numpy>=2.2.5
+Dynamic: license-file
+
+# Batch Tensorsolve
+
+<p align="center">
+  <a href="https://github.com/34j/batch-tensorsolve/actions/workflows/ci.yml?query=branch%3Amain">
+    <img src="https://img.shields.io/github/actions/workflow/status/34j/batch-tensorsolve/ci.yml?branch=main&label=CI&logo=github&style=flat-square" alt="CI Status" >
+  </a>
+  <a href="https://batch-tensorsolve.readthedocs.io">
+    <img src="https://img.shields.io/readthedocs/batch-tensorsolve.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
+  </a>
+  <a href="https://codecov.io/gh/34j/batch-tensorsolve">
+    <img src="https://img.shields.io/codecov/c/github/34j/batch-tensorsolve.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
+  </a>
+</p>
+<p align="center">
+  <a href="https://github.com/astral-sh/uv">
+    <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
+  </a>
+  <a href="https://github.com/astral-sh/ruff">
+    <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
+  </a>
+  <a href="https://github.com/pre-commit/pre-commit">
+    <img src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white&style=flat-square" alt="pre-commit">
+  </a>
+</p>
+<p align="center">
+  <a href="https://pypi.org/project/batch-tensorsolve/">
+    <img src="https://img.shields.io/pypi/v/batch-tensorsolve.svg?logo=python&logoColor=fff&style=flat-square" alt="PyPI Version">
+  </a>
+  <img src="https://img.shields.io/pypi/pyversions/batch-tensorsolve.svg?style=flat-square&logo=python&amp;logoColor=fff" alt="Supported Python versions">
+  <img src="https://img.shields.io/pypi/l/batch-tensorsolve.svg?style=flat-square" alt="License">
+</p>
+
+---
+
+**Documentation**: <a href="https://batch-tensorsolve.readthedocs.io" target="_blank">https://batch-tensorsolve.readthedocs.io </a>
+
+**Source Code**: <a href="https://github.com/34j/batch-tensorsolve" target="_blank">https://github.com/34j/batch-tensorsolve </a>
+
+---
+
+Batched tensorsolve() for NumPy / PyTorch / JAX. ([numpy/numpy#28099](https://github.com/numpy/numpy/issues/28099))
+
+## Installation
+
+Install this via pip (or your favourite package manager):
+
+```shell
+pip install batch-tensorsolve
+```
+
+## Usage
+
+```python
+import numpy as np
+from numpy.testing import assert_allclose
+
+from batch_tensorsolve import btensorsolve
+
+a = np.random.randn(2, 2, 3, 6)
+b = np.random.randn(2, 2, 3)
+assert_allclose(np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b)), b)
+```
+
+## Advanced Usage
+
+It is recommended to explicitly specify the batch axes, as the desired shape will be ambiguous if axes of size 1 are present.
+
+```python
+import numpy as np
+
+from batch_tensorsolve import btensorsolve
+
+a = np.random.randn(2, 1, 2, 2)
+b = np.random.randn(2, 1, 2)
+# 2 possibilities:
+assert btensorsolve(a, b, num_batch_axes=1).shape == (2, 2) # 1st axis is batch
+assert btensorsolve(a, b, num_batch_axes=2).shape == (2, 1, 2) # 1st and 2nd axes are batch
+```
+
+Broadcasting-like behavior is also supported:
+
+```python
+import numpy as np
+from numpy.testing import assert_allclose
+
+from batch_tensorsolve import btensorsolve
+
+a = np.random.randn(1, 2, 3, 6) # -> (2, 2, 3, 6)
+b = np.random.randn(2, 1, 1) # -> (2, 2, 3)
+left = np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b))
+assert_allclose(left, np.broadcast_to(b, left.shape))
+```
+
+Note that broadcasting (repeating) `a` for non-batch axes will result in `numpy.linalg.LinAlgError: Singular matrix` because the matrix representation of `a` has duplicate rows.
+
+## Contributors ✨
+
+Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
+
+<!-- prettier-ignore-start -->
+<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
+<!-- markdownlint-disable -->
+<!-- markdownlint-enable -->
+<!-- ALL-CONTRIBUTORS-LIST:END -->
+<!-- prettier-ignore-end -->
+
+This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
+
+## Credits
+
+[![Copier](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/copier-org/copier/master/img/badge/badge-grayscale-inverted-border-orange.json)](https://github.com/copier-org/copier)
+
+This package was created with
+[Copier](https://copier.readthedocs.io/) and the
+[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
+project template.

batch_tensorsolve-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+batch_tensorsolve/__init__.py,sha256=9ggnkEF1DPkJqGEZERfUnpjmWS8sU6QBSbRk79Lvbdg,198
+batch_tensorsolve/_main.py,sha256=u7pypXQumqMXq_mANq3gxGGO0iPQp54zQ6mnW1gYOIs,5677
+batch_tensorsolve/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+batch_tensorsolve-0.1.0.dist-info/licenses/LICENSE,sha256=Otwb96xcK4QBqXXVdBKpLZVnc4mJuFTa8rsMUEpphMM,1061
+batch_tensorsolve-0.1.0.dist-info/METADATA,sha256=HX5cZa2JiQOeVEYBeHrGDsk8lG1o5h5QnCptoxSeABM,5631
+batch_tensorsolve-0.1.0.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+batch_tensorsolve-0.1.0.dist-info/top_level.txt,sha256=C2-1orVxMMM7moCnIWJ35hk_aPGfh7P1ulR1hExyaUU,18
+batch_tensorsolve-0.1.0.dist-info/RECORD,,

batch_tensorsolve-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.4.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

batch_tensorsolve-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,22 @@
+
+MIT License
+
+Copyright (c) 2025 34j
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

batch_tensorsolve-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+batch_tensorsolve