batch-tensorsolve 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- batch_tensorsolve-0.1.0/LICENSE +22 -0
- batch_tensorsolve-0.1.0/PKG-INFO +142 -0
- batch_tensorsolve-0.1.0/README.md +116 -0
- batch_tensorsolve-0.1.0/pyproject.toml +161 -0
- batch_tensorsolve-0.1.0/setup.cfg +4 -0
- batch_tensorsolve-0.1.0/setup.py +9 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve/__init__.py +4 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve/_main.py +178 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve/py.typed +0 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve.egg-info/PKG-INFO +142 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve.egg-info/SOURCES.txt +13 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve.egg-info/dependency_links.txt +1 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve.egg-info/requires.txt +2 -0
- batch_tensorsolve-0.1.0/src/batch_tensorsolve.egg-info/top_level.txt +1 -0
- batch_tensorsolve-0.1.0/tests/test_main.py +37 -0
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
|
|
2
|
+
MIT License
|
|
3
|
+
|
|
4
|
+
Copyright (c) 2025 34j
|
|
5
|
+
|
|
6
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
7
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
8
|
+
in the Software without restriction, including without limitation the rights
|
|
9
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
10
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
11
|
+
furnished to do so, subject to the following conditions:
|
|
12
|
+
|
|
13
|
+
The above copyright notice and this permission notice shall be included in all
|
|
14
|
+
copies or substantial portions of the Software.
|
|
15
|
+
|
|
16
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
17
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
18
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
19
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
20
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
21
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
22
|
+
SOFTWARE.
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: batch-tensorsolve
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Batched tensorsolve() for NumPy / PyTorch / JAX
|
|
5
|
+
Author-email: 34j <34j.95a2p@simplelogin.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Bug Tracker, https://github.com/34j/batch-tensorsolve/issues
|
|
8
|
+
Project-URL: Changelog, https://github.com/34j/batch-tensorsolve/blob/main/CHANGELOG.md
|
|
9
|
+
Project-URL: documentation, https://batch-tensorsolve.readthedocs.io
|
|
10
|
+
Project-URL: repository, https://github.com/34j/batch-tensorsolve
|
|
11
|
+
Classifier: Development Status :: 2 - Pre-Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Natural Language :: English
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: array-api-compat>=1.11.2
|
|
24
|
+
Requires-Dist: numpy>=2.2.5
|
|
25
|
+
Dynamic: license-file
|
|
26
|
+
|
|
27
|
+
# Batch Tensorsolve
|
|
28
|
+
|
|
29
|
+
<p align="center">
|
|
30
|
+
<a href="https://github.com/34j/batch-tensorsolve/actions/workflows/ci.yml?query=branch%3Amain">
|
|
31
|
+
<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" >
|
|
32
|
+
</a>
|
|
33
|
+
<a href="https://batch-tensorsolve.readthedocs.io">
|
|
34
|
+
<img src="https://img.shields.io/readthedocs/batch-tensorsolve.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
35
|
+
</a>
|
|
36
|
+
<a href="https://codecov.io/gh/34j/batch-tensorsolve">
|
|
37
|
+
<img src="https://img.shields.io/codecov/c/github/34j/batch-tensorsolve.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
38
|
+
</a>
|
|
39
|
+
</p>
|
|
40
|
+
<p align="center">
|
|
41
|
+
<a href="https://github.com/astral-sh/uv">
|
|
42
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
43
|
+
</a>
|
|
44
|
+
<a href="https://github.com/astral-sh/ruff">
|
|
45
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
46
|
+
</a>
|
|
47
|
+
<a href="https://github.com/pre-commit/pre-commit">
|
|
48
|
+
<img src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white&style=flat-square" alt="pre-commit">
|
|
49
|
+
</a>
|
|
50
|
+
</p>
|
|
51
|
+
<p align="center">
|
|
52
|
+
<a href="https://pypi.org/project/batch-tensorsolve/">
|
|
53
|
+
<img src="https://img.shields.io/pypi/v/batch-tensorsolve.svg?logo=python&logoColor=fff&style=flat-square" alt="PyPI Version">
|
|
54
|
+
</a>
|
|
55
|
+
<img src="https://img.shields.io/pypi/pyversions/batch-tensorsolve.svg?style=flat-square&logo=python&logoColor=fff" alt="Supported Python versions">
|
|
56
|
+
<img src="https://img.shields.io/pypi/l/batch-tensorsolve.svg?style=flat-square" alt="License">
|
|
57
|
+
</p>
|
|
58
|
+
|
|
59
|
+
---
|
|
60
|
+
|
|
61
|
+
**Documentation**: <a href="https://batch-tensorsolve.readthedocs.io" target="_blank">https://batch-tensorsolve.readthedocs.io </a>
|
|
62
|
+
|
|
63
|
+
**Source Code**: <a href="https://github.com/34j/batch-tensorsolve" target="_blank">https://github.com/34j/batch-tensorsolve </a>
|
|
64
|
+
|
|
65
|
+
---
|
|
66
|
+
|
|
67
|
+
Batched tensorsolve() for NumPy / PyTorch / JAX. ([numpy/numpy#28099](https://github.com/numpy/numpy/issues/28099))
|
|
68
|
+
|
|
69
|
+
## Installation
|
|
70
|
+
|
|
71
|
+
Install this via pip (or your favourite package manager):
|
|
72
|
+
|
|
73
|
+
```shell
|
|
74
|
+
pip install batch-tensorsolve
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
## Usage
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
import numpy as np
|
|
81
|
+
from numpy.testing import assert_allclose
|
|
82
|
+
|
|
83
|
+
from batch_tensorsolve import btensorsolve
|
|
84
|
+
|
|
85
|
+
a = np.random.randn(2, 2, 3, 6)
|
|
86
|
+
b = np.random.randn(2, 2, 3)
|
|
87
|
+
assert_allclose(np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b)), b)
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Advanced Usage
|
|
91
|
+
|
|
92
|
+
It is recommended to explicitly specify the batch axes, as the desired shape will be ambiguous if axes of size 1 are present.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
import numpy as np
|
|
96
|
+
|
|
97
|
+
from batch_tensorsolve import btensorsolve
|
|
98
|
+
|
|
99
|
+
a = np.random.randn(2, 1, 2, 2)
|
|
100
|
+
b = np.random.randn(2, 1, 2)
|
|
101
|
+
# 2 possibilities:
|
|
102
|
+
assert btensorsolve(a, b, num_batch_axes=1).shape == (2, 2) # 1st axis is batch
|
|
103
|
+
assert btensorsolve(a, b, num_batch_axes=2).shape == (2, 1, 2) # 1st and 2nd axes are batch
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Broadcasting-like behavior is also supported:
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
import numpy as np
|
|
110
|
+
from numpy.testing import assert_allclose
|
|
111
|
+
|
|
112
|
+
from batch_tensorsolve import btensorsolve
|
|
113
|
+
|
|
114
|
+
a = np.random.randn(1, 2, 3, 6) # -> (2, 2, 3, 6)
|
|
115
|
+
b = np.random.randn(2, 1, 1) # -> (2, 2, 3)
|
|
116
|
+
left = np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b))
|
|
117
|
+
assert_allclose(left, np.broadcast_to(b, left.shape))
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
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.
|
|
121
|
+
|
|
122
|
+
## Contributors ✨
|
|
123
|
+
|
|
124
|
+
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
125
|
+
|
|
126
|
+
<!-- prettier-ignore-start -->
|
|
127
|
+
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
128
|
+
<!-- markdownlint-disable -->
|
|
129
|
+
<!-- markdownlint-enable -->
|
|
130
|
+
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
131
|
+
<!-- prettier-ignore-end -->
|
|
132
|
+
|
|
133
|
+
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
134
|
+
|
|
135
|
+
## Credits
|
|
136
|
+
|
|
137
|
+
[](https://github.com/copier-org/copier)
|
|
138
|
+
|
|
139
|
+
This package was created with
|
|
140
|
+
[Copier](https://copier.readthedocs.io/) and the
|
|
141
|
+
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
142
|
+
project template.
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
# Batch Tensorsolve
|
|
2
|
+
|
|
3
|
+
<p align="center">
|
|
4
|
+
<a href="https://github.com/34j/batch-tensorsolve/actions/workflows/ci.yml?query=branch%3Amain">
|
|
5
|
+
<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" >
|
|
6
|
+
</a>
|
|
7
|
+
<a href="https://batch-tensorsolve.readthedocs.io">
|
|
8
|
+
<img src="https://img.shields.io/readthedocs/batch-tensorsolve.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
9
|
+
</a>
|
|
10
|
+
<a href="https://codecov.io/gh/34j/batch-tensorsolve">
|
|
11
|
+
<img src="https://img.shields.io/codecov/c/github/34j/batch-tensorsolve.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
12
|
+
</a>
|
|
13
|
+
</p>
|
|
14
|
+
<p align="center">
|
|
15
|
+
<a href="https://github.com/astral-sh/uv">
|
|
16
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
17
|
+
</a>
|
|
18
|
+
<a href="https://github.com/astral-sh/ruff">
|
|
19
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
20
|
+
</a>
|
|
21
|
+
<a href="https://github.com/pre-commit/pre-commit">
|
|
22
|
+
<img src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white&style=flat-square" alt="pre-commit">
|
|
23
|
+
</a>
|
|
24
|
+
</p>
|
|
25
|
+
<p align="center">
|
|
26
|
+
<a href="https://pypi.org/project/batch-tensorsolve/">
|
|
27
|
+
<img src="https://img.shields.io/pypi/v/batch-tensorsolve.svg?logo=python&logoColor=fff&style=flat-square" alt="PyPI Version">
|
|
28
|
+
</a>
|
|
29
|
+
<img src="https://img.shields.io/pypi/pyversions/batch-tensorsolve.svg?style=flat-square&logo=python&logoColor=fff" alt="Supported Python versions">
|
|
30
|
+
<img src="https://img.shields.io/pypi/l/batch-tensorsolve.svg?style=flat-square" alt="License">
|
|
31
|
+
</p>
|
|
32
|
+
|
|
33
|
+
---
|
|
34
|
+
|
|
35
|
+
**Documentation**: <a href="https://batch-tensorsolve.readthedocs.io" target="_blank">https://batch-tensorsolve.readthedocs.io </a>
|
|
36
|
+
|
|
37
|
+
**Source Code**: <a href="https://github.com/34j/batch-tensorsolve" target="_blank">https://github.com/34j/batch-tensorsolve </a>
|
|
38
|
+
|
|
39
|
+
---
|
|
40
|
+
|
|
41
|
+
Batched tensorsolve() for NumPy / PyTorch / JAX. ([numpy/numpy#28099](https://github.com/numpy/numpy/issues/28099))
|
|
42
|
+
|
|
43
|
+
## Installation
|
|
44
|
+
|
|
45
|
+
Install this via pip (or your favourite package manager):
|
|
46
|
+
|
|
47
|
+
```shell
|
|
48
|
+
pip install batch-tensorsolve
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
## Usage
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
import numpy as np
|
|
55
|
+
from numpy.testing import assert_allclose
|
|
56
|
+
|
|
57
|
+
from batch_tensorsolve import btensorsolve
|
|
58
|
+
|
|
59
|
+
a = np.random.randn(2, 2, 3, 6)
|
|
60
|
+
b = np.random.randn(2, 2, 3)
|
|
61
|
+
assert_allclose(np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b)), b)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## Advanced Usage
|
|
65
|
+
|
|
66
|
+
It is recommended to explicitly specify the batch axes, as the desired shape will be ambiguous if axes of size 1 are present.
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
import numpy as np
|
|
70
|
+
|
|
71
|
+
from batch_tensorsolve import btensorsolve
|
|
72
|
+
|
|
73
|
+
a = np.random.randn(2, 1, 2, 2)
|
|
74
|
+
b = np.random.randn(2, 1, 2)
|
|
75
|
+
# 2 possibilities:
|
|
76
|
+
assert btensorsolve(a, b, num_batch_axes=1).shape == (2, 2) # 1st axis is batch
|
|
77
|
+
assert btensorsolve(a, b, num_batch_axes=2).shape == (2, 1, 2) # 1st and 2nd axes are batch
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
Broadcasting-like behavior is also supported:
|
|
81
|
+
|
|
82
|
+
```python
|
|
83
|
+
import numpy as np
|
|
84
|
+
from numpy.testing import assert_allclose
|
|
85
|
+
|
|
86
|
+
from batch_tensorsolve import btensorsolve
|
|
87
|
+
|
|
88
|
+
a = np.random.randn(1, 2, 3, 6) # -> (2, 2, 3, 6)
|
|
89
|
+
b = np.random.randn(2, 1, 1) # -> (2, 2, 3)
|
|
90
|
+
left = np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b))
|
|
91
|
+
assert_allclose(left, np.broadcast_to(b, left.shape))
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
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.
|
|
95
|
+
|
|
96
|
+
## Contributors ✨
|
|
97
|
+
|
|
98
|
+
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
99
|
+
|
|
100
|
+
<!-- prettier-ignore-start -->
|
|
101
|
+
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
102
|
+
<!-- markdownlint-disable -->
|
|
103
|
+
<!-- markdownlint-enable -->
|
|
104
|
+
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
105
|
+
<!-- prettier-ignore-end -->
|
|
106
|
+
|
|
107
|
+
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
108
|
+
|
|
109
|
+
## Credits
|
|
110
|
+
|
|
111
|
+
[](https://github.com/copier-org/copier)
|
|
112
|
+
|
|
113
|
+
This package was created with
|
|
114
|
+
[Copier](https://copier.readthedocs.io/) and the
|
|
115
|
+
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
116
|
+
project template.
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
build-backend = "setuptools.build_meta"
|
|
3
|
+
requires = [ "setuptools" ]
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "batch-tensorsolve"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Batched tensorsolve() for NumPy / PyTorch / JAX"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "MIT" }
|
|
11
|
+
authors = [
|
|
12
|
+
{ name = "34j", email = "34j.95a2p@simplelogin.com" },
|
|
13
|
+
]
|
|
14
|
+
requires-python = ">=3.10"
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Development Status :: 2 - Pre-Alpha",
|
|
17
|
+
"Intended Audience :: Developers",
|
|
18
|
+
"Natural Language :: English",
|
|
19
|
+
"Operating System :: OS Independent",
|
|
20
|
+
"Programming Language :: Python :: 3.10",
|
|
21
|
+
"Programming Language :: Python :: 3.11",
|
|
22
|
+
"Programming Language :: Python :: 3.12",
|
|
23
|
+
"Programming Language :: Python :: 3.13",
|
|
24
|
+
"Topic :: Software Development :: Libraries",
|
|
25
|
+
]
|
|
26
|
+
|
|
27
|
+
dependencies = [
|
|
28
|
+
"array-api-compat>=1.11.2",
|
|
29
|
+
"numpy>=2.2.5",
|
|
30
|
+
]
|
|
31
|
+
urls."Bug Tracker" = "https://github.com/34j/batch-tensorsolve/issues"
|
|
32
|
+
urls.Changelog = "https://github.com/34j/batch-tensorsolve/blob/main/CHANGELOG.md"
|
|
33
|
+
urls.documentation = "https://batch-tensorsolve.readthedocs.io"
|
|
34
|
+
urls.repository = "https://github.com/34j/batch-tensorsolve"
|
|
35
|
+
|
|
36
|
+
[dependency-groups]
|
|
37
|
+
dev = [
|
|
38
|
+
"pytest>=8,<9",
|
|
39
|
+
"pytest-cov>=6,<7",
|
|
40
|
+
]
|
|
41
|
+
docs = [
|
|
42
|
+
"furo>=2023.5.20; python_version>='3.11'",
|
|
43
|
+
"myst-parser>=0.16; python_version>='3.11'",
|
|
44
|
+
"sphinx>=4; python_version>='3.11'",
|
|
45
|
+
"sphinx-autobuild>=2024,<2025; python_version>='3.11'",
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
[tool.ruff]
|
|
49
|
+
line-length = 88
|
|
50
|
+
lint.select = [
|
|
51
|
+
"B", # flake8-bugbear
|
|
52
|
+
"D", # flake8-docstrings
|
|
53
|
+
"C4", # flake8-comprehensions
|
|
54
|
+
"S", # flake8-bandit
|
|
55
|
+
"F", # pyflake
|
|
56
|
+
"E", # pycodestyle
|
|
57
|
+
"W", # pycodestyle
|
|
58
|
+
"UP", # pyupgrade
|
|
59
|
+
"I", # isort
|
|
60
|
+
"RUF", # ruff specific
|
|
61
|
+
]
|
|
62
|
+
lint.ignore = [
|
|
63
|
+
"D203", # 1 blank line required before class docstring
|
|
64
|
+
"D212", # Multi-line docstring summary should start at the first line
|
|
65
|
+
"D100", # Missing docstring in public module
|
|
66
|
+
"D104", # Missing docstring in public package
|
|
67
|
+
"D107", # Missing docstring in `__init__`
|
|
68
|
+
"D401", # First line of docstring should be in imperative mood
|
|
69
|
+
]
|
|
70
|
+
lint.per-file-ignores."conftest.py" = [ "D100" ]
|
|
71
|
+
lint.per-file-ignores."docs/conf.py" = [ "D100" ]
|
|
72
|
+
lint.per-file-ignores."setup.py" = [ "D100" ]
|
|
73
|
+
lint.per-file-ignores."tests/**/*" = [
|
|
74
|
+
"D100",
|
|
75
|
+
"D101",
|
|
76
|
+
"D102",
|
|
77
|
+
"D103",
|
|
78
|
+
"D104",
|
|
79
|
+
"S101",
|
|
80
|
+
]
|
|
81
|
+
lint.isort.known-first-party = [ "batch_tensorsolve", "tests" ]
|
|
82
|
+
|
|
83
|
+
[tool.pytest.ini_options]
|
|
84
|
+
addopts = """\
|
|
85
|
+
-v
|
|
86
|
+
-Wdefault
|
|
87
|
+
--cov=batch_tensorsolve
|
|
88
|
+
--cov-report=term
|
|
89
|
+
--cov-report=xml
|
|
90
|
+
"""
|
|
91
|
+
pythonpath = [ "src" ]
|
|
92
|
+
|
|
93
|
+
[tool.coverage.run]
|
|
94
|
+
branch = true
|
|
95
|
+
|
|
96
|
+
[tool.coverage.report]
|
|
97
|
+
exclude_lines = [
|
|
98
|
+
"pragma: no cover",
|
|
99
|
+
"@overload",
|
|
100
|
+
"if TYPE_CHECKING",
|
|
101
|
+
"raise NotImplementedError",
|
|
102
|
+
'if __name__ == "__main__":',
|
|
103
|
+
]
|
|
104
|
+
|
|
105
|
+
[tool.mypy]
|
|
106
|
+
check_untyped_defs = true
|
|
107
|
+
disallow_any_generics = true
|
|
108
|
+
disallow_incomplete_defs = true
|
|
109
|
+
disallow_untyped_defs = true
|
|
110
|
+
mypy_path = "src/"
|
|
111
|
+
no_implicit_optional = true
|
|
112
|
+
show_error_codes = true
|
|
113
|
+
warn_unreachable = true
|
|
114
|
+
warn_unused_ignores = true
|
|
115
|
+
exclude = [
|
|
116
|
+
'docs/.*',
|
|
117
|
+
'setup.py',
|
|
118
|
+
]
|
|
119
|
+
|
|
120
|
+
[[tool.mypy.overrides]]
|
|
121
|
+
module = "tests.*"
|
|
122
|
+
allow_untyped_defs = true
|
|
123
|
+
|
|
124
|
+
[[tool.mypy.overrides]]
|
|
125
|
+
module = "docs.*"
|
|
126
|
+
ignore_errors = true
|
|
127
|
+
|
|
128
|
+
[tool.semantic_release]
|
|
129
|
+
version_toml = [ "pyproject.toml:project.version" ]
|
|
130
|
+
version_variables = [
|
|
131
|
+
"src/batch_tensorsolve/__init__.py:__version__",
|
|
132
|
+
"docs/conf.py:release",
|
|
133
|
+
]
|
|
134
|
+
build_command = """
|
|
135
|
+
pip install uv
|
|
136
|
+
uv lock
|
|
137
|
+
git add uv.lock
|
|
138
|
+
uv build
|
|
139
|
+
"""
|
|
140
|
+
|
|
141
|
+
[tool.semantic_release.changelog]
|
|
142
|
+
exclude_commit_patterns = [
|
|
143
|
+
'''chore(?:\([^)]*?\))?: .+''',
|
|
144
|
+
'''ci(?:\([^)]*?\))?: .+''',
|
|
145
|
+
'''refactor(?:\([^)]*?\))?: .+''',
|
|
146
|
+
'''style(?:\([^)]*?\))?: .+''',
|
|
147
|
+
'''test(?:\([^)]*?\))?: .+''',
|
|
148
|
+
'''build\((?!deps\): .+)''',
|
|
149
|
+
'''Merged? .*''',
|
|
150
|
+
'''Initial [Cc]ommit.*''', # codespell:ignore
|
|
151
|
+
]
|
|
152
|
+
|
|
153
|
+
[tool.semantic_release.changelog.environment]
|
|
154
|
+
keep_trailing_newline = true
|
|
155
|
+
|
|
156
|
+
[tool.semantic_release.branches.main]
|
|
157
|
+
match = "main"
|
|
158
|
+
|
|
159
|
+
[tool.semantic_release.branches.noop]
|
|
160
|
+
match = "(?!main$)"
|
|
161
|
+
prerelease = true
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
import warnings
|
|
2
|
+
from typing import TypeVar
|
|
3
|
+
|
|
4
|
+
from array_api_compat import array_namespace
|
|
5
|
+
|
|
6
|
+
TArray = TypeVar("TArray")
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class AmbiguousBatchAxesWarning(RuntimeWarning):
|
|
10
|
+
pass
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def broadcast_without_repeating(
|
|
14
|
+
*arrays: TArray, check_same_ndim: bool = False
|
|
15
|
+
) -> tuple[TArray, ...]:
|
|
16
|
+
"""
|
|
17
|
+
Broadcast arrays without repeating the data.
|
|
18
|
+
|
|
19
|
+
Parameters
|
|
20
|
+
----------
|
|
21
|
+
arrays : TArray
|
|
22
|
+
The arrays to broadcast.
|
|
23
|
+
check_same_ndim : bool, optional
|
|
24
|
+
Whether to check if all arrays have the same number of dimensions.
|
|
25
|
+
Default is False.
|
|
26
|
+
|
|
27
|
+
Returns
|
|
28
|
+
-------
|
|
29
|
+
tuple[TArray]
|
|
30
|
+
The broadcasted arrays.
|
|
31
|
+
|
|
32
|
+
"""
|
|
33
|
+
xp = array_namespace(*arrays)
|
|
34
|
+
arrays_ = tuple(xp.asarray(a) for a in arrays)
|
|
35
|
+
xp.broadcast_shapes(*[a.shape for a in arrays_])
|
|
36
|
+
if check_same_ndim:
|
|
37
|
+
if len({a.ndim for a in arrays_}) != 1:
|
|
38
|
+
raise ValueError(
|
|
39
|
+
"All arrays must have the same number of dimensions, "
|
|
40
|
+
f"but got {tuple(a.ndim for a in arrays_)}"
|
|
41
|
+
)
|
|
42
|
+
return arrays_
|
|
43
|
+
max_dim = max(a.ndim for a in arrays_)
|
|
44
|
+
return tuple(array[(None,) * (max_dim - array.ndim) + (...,)] for array in arrays_)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def btensorsolve(
|
|
48
|
+
a: TArray,
|
|
49
|
+
b: TArray,
|
|
50
|
+
/,
|
|
51
|
+
*,
|
|
52
|
+
axes: tuple[int] | None = None,
|
|
53
|
+
num_batch_axes: int | None = None,
|
|
54
|
+
) -> TArray:
|
|
55
|
+
"""
|
|
56
|
+
Solve the tensor equation ``a x = b`` for x.
|
|
57
|
+
|
|
58
|
+
It is assumed that all indices of `x` are summed over in the product,
|
|
59
|
+
together with the rightmost indices of `a`, as is done in, for example,
|
|
60
|
+
``tensordot(a, x, axes=x.ndim)``.
|
|
61
|
+
|
|
62
|
+
Parameters
|
|
63
|
+
----------
|
|
64
|
+
a : array_like
|
|
65
|
+
Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
|
|
66
|
+
the shape of that sub-tensor of `a` consisting of the appropriate
|
|
67
|
+
number of its rightmost indices, and must be such that
|
|
68
|
+
**there exists i that**
|
|
69
|
+
``prod(Q) == prod(b.shape[i:])``
|
|
70
|
+
b : array_like
|
|
71
|
+
Right-hand tensor, which can be of any shape.
|
|
72
|
+
axes : tuple of ints, optional
|
|
73
|
+
Axes in `a` to reorder to the right, before inversion.
|
|
74
|
+
If None (default), no reordering is done.
|
|
75
|
+
num_batch_axes : int, optional
|
|
76
|
+
The number of batch dimensions. If None (default), the number of batch
|
|
77
|
+
dimensions is inferred from the shapes of `a` and `b`.
|
|
78
|
+
|
|
79
|
+
Let ``shape = np.broadcast_shapes(a.shape[:b.ndim],
|
|
80
|
+
b.shape) + a.shape[b.ndim:]``.
|
|
81
|
+
|
|
82
|
+
It is recommended to specify this argument, as the inference
|
|
83
|
+
might be wrong if there exists i >= 0, j > 0 that
|
|
84
|
+
``prod(shape[:i]) == prod(shape[i+j:])``. (j + 1 possibilities)
|
|
85
|
+
|
|
86
|
+
For example, if `a` has shape (3, 1, 1, 2, 2) and `b` has shape
|
|
87
|
+
(3, 1, 1, 2), it is possible that
|
|
88
|
+
- axis 1 is the batch axes and desired output shape is (3, 2)
|
|
89
|
+
- axis 1, 2 are the batch axes and desired output shape is (3, 1, 2)
|
|
90
|
+
- axis 1, 2, 3 are the batch axes and desired output shape is (3, 1, 1, 2)
|
|
91
|
+
|
|
92
|
+
Returns
|
|
93
|
+
-------
|
|
94
|
+
x : ndarray, shape Q
|
|
95
|
+
|
|
96
|
+
Raises
|
|
97
|
+
------
|
|
98
|
+
LinAlgError
|
|
99
|
+
If `a` is singular or not 'square' (in the above sense).
|
|
100
|
+
|
|
101
|
+
Warnings
|
|
102
|
+
--------
|
|
103
|
+
AmbiguousBatchAxesWarning
|
|
104
|
+
If the number of batch axes cannot be inferred from the shapes of
|
|
105
|
+
`a` and `b`, and `num_batch_axes` is not specified.
|
|
106
|
+
|
|
107
|
+
See Also
|
|
108
|
+
--------
|
|
109
|
+
numpy.tensordot, tensorinv, numpy.einsum
|
|
110
|
+
|
|
111
|
+
Examples
|
|
112
|
+
--------
|
|
113
|
+
>>> import numpy as np
|
|
114
|
+
>>> rng = np.random.default_rng()
|
|
115
|
+
>>> a = rng.normal(size=(2, 2*3, 4, 2, 3, 4))
|
|
116
|
+
>>> b = rng.normal(size=(2, 2*3, 4))
|
|
117
|
+
>>> x = np.linalg.tensorsolve(a, b)
|
|
118
|
+
>>> x.shape
|
|
119
|
+
(2, 2, 3, 4)
|
|
120
|
+
>>> np.allclose(np.einsum('...ijklm,...klm->...ij', a, x), b)
|
|
121
|
+
True
|
|
122
|
+
|
|
123
|
+
"""
|
|
124
|
+
# https://github.com/numpy/numpy/blob/
|
|
125
|
+
# e7a123b2d3eca9897843791dd698c1803d9a39c2/numpy/linalg/_linalg.py#L291
|
|
126
|
+
xp = array_namespace(a, b)
|
|
127
|
+
a_ = xp.asarray(a)
|
|
128
|
+
b_ = xp.asarray(b)
|
|
129
|
+
an = a_.ndim
|
|
130
|
+
if axes is not None:
|
|
131
|
+
allaxes = list(range(0, an))
|
|
132
|
+
for k in axes:
|
|
133
|
+
allaxes.remove(k)
|
|
134
|
+
allaxes.insert(an, k)
|
|
135
|
+
a_ = a_.transpose(allaxes)
|
|
136
|
+
|
|
137
|
+
# find right dimensions
|
|
138
|
+
# a = [2 (dim1) 2 2 3 (dim2) 2 6]
|
|
139
|
+
# b = [2 (dim1) 2 2 3 (dim2)]
|
|
140
|
+
axis_sol_last = b_.ndim
|
|
141
|
+
ashape = a_.shape
|
|
142
|
+
bshape = b_.shape
|
|
143
|
+
shape = xp.broadcast_shapes(ashape[:axis_sol_last], bshape)
|
|
144
|
+
|
|
145
|
+
# the dimension of the linear system
|
|
146
|
+
sol_size = int(xp.prod(ashape[axis_sol_last:]))
|
|
147
|
+
if num_batch_axes is None:
|
|
148
|
+
sol_size_current = 1
|
|
149
|
+
for num_batch_axes in range(axis_sol_last - 1, -1, -1):
|
|
150
|
+
sol_size_current *= shape[num_batch_axes]
|
|
151
|
+
if sol_size_current == sol_size:
|
|
152
|
+
break
|
|
153
|
+
else:
|
|
154
|
+
raise ValueError(
|
|
155
|
+
"Unable to divide batch dimensions and solution dimensions"
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
if num_batch_axes > 0 and shape[num_batch_axes - 1] == 1:
|
|
159
|
+
warnings.warn(
|
|
160
|
+
"It is impossible to infer the number of "
|
|
161
|
+
"batch axes from the shapes of `a` and `b`. "
|
|
162
|
+
"Consider specifying `num_batch_axes` explicitly.",
|
|
163
|
+
AmbiguousBatchAxesWarning,
|
|
164
|
+
stacklevel=2,
|
|
165
|
+
)
|
|
166
|
+
|
|
167
|
+
a_ = xp.broadcast_to(
|
|
168
|
+
a_,
|
|
169
|
+
ashape[:num_batch_axes]
|
|
170
|
+
+ shape[num_batch_axes:axis_sol_last]
|
|
171
|
+
+ ashape[axis_sol_last:],
|
|
172
|
+
)
|
|
173
|
+
b_ = xp.broadcast_to(b_, bshape[:num_batch_axes] + shape[num_batch_axes:])
|
|
174
|
+
a_ = a_.reshape(ashape[:num_batch_axes] + (sol_size, sol_size))
|
|
175
|
+
b_ = b_.reshape(bshape[:num_batch_axes] + (sol_size, 1))
|
|
176
|
+
x = xp.linalg.solve(a_, b_)
|
|
177
|
+
x = x.reshape(shape[:num_batch_axes] + ashape[axis_sol_last:])
|
|
178
|
+
return x
|
|
File without changes
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: batch-tensorsolve
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Batched tensorsolve() for NumPy / PyTorch / JAX
|
|
5
|
+
Author-email: 34j <34j.95a2p@simplelogin.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Bug Tracker, https://github.com/34j/batch-tensorsolve/issues
|
|
8
|
+
Project-URL: Changelog, https://github.com/34j/batch-tensorsolve/blob/main/CHANGELOG.md
|
|
9
|
+
Project-URL: documentation, https://batch-tensorsolve.readthedocs.io
|
|
10
|
+
Project-URL: repository, https://github.com/34j/batch-tensorsolve
|
|
11
|
+
Classifier: Development Status :: 2 - Pre-Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Natural Language :: English
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: array-api-compat>=1.11.2
|
|
24
|
+
Requires-Dist: numpy>=2.2.5
|
|
25
|
+
Dynamic: license-file
|
|
26
|
+
|
|
27
|
+
# Batch Tensorsolve
|
|
28
|
+
|
|
29
|
+
<p align="center">
|
|
30
|
+
<a href="https://github.com/34j/batch-tensorsolve/actions/workflows/ci.yml?query=branch%3Amain">
|
|
31
|
+
<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" >
|
|
32
|
+
</a>
|
|
33
|
+
<a href="https://batch-tensorsolve.readthedocs.io">
|
|
34
|
+
<img src="https://img.shields.io/readthedocs/batch-tensorsolve.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
35
|
+
</a>
|
|
36
|
+
<a href="https://codecov.io/gh/34j/batch-tensorsolve">
|
|
37
|
+
<img src="https://img.shields.io/codecov/c/github/34j/batch-tensorsolve.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
38
|
+
</a>
|
|
39
|
+
</p>
|
|
40
|
+
<p align="center">
|
|
41
|
+
<a href="https://github.com/astral-sh/uv">
|
|
42
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
43
|
+
</a>
|
|
44
|
+
<a href="https://github.com/astral-sh/ruff">
|
|
45
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
46
|
+
</a>
|
|
47
|
+
<a href="https://github.com/pre-commit/pre-commit">
|
|
48
|
+
<img src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white&style=flat-square" alt="pre-commit">
|
|
49
|
+
</a>
|
|
50
|
+
</p>
|
|
51
|
+
<p align="center">
|
|
52
|
+
<a href="https://pypi.org/project/batch-tensorsolve/">
|
|
53
|
+
<img src="https://img.shields.io/pypi/v/batch-tensorsolve.svg?logo=python&logoColor=fff&style=flat-square" alt="PyPI Version">
|
|
54
|
+
</a>
|
|
55
|
+
<img src="https://img.shields.io/pypi/pyversions/batch-tensorsolve.svg?style=flat-square&logo=python&logoColor=fff" alt="Supported Python versions">
|
|
56
|
+
<img src="https://img.shields.io/pypi/l/batch-tensorsolve.svg?style=flat-square" alt="License">
|
|
57
|
+
</p>
|
|
58
|
+
|
|
59
|
+
---
|
|
60
|
+
|
|
61
|
+
**Documentation**: <a href="https://batch-tensorsolve.readthedocs.io" target="_blank">https://batch-tensorsolve.readthedocs.io </a>
|
|
62
|
+
|
|
63
|
+
**Source Code**: <a href="https://github.com/34j/batch-tensorsolve" target="_blank">https://github.com/34j/batch-tensorsolve </a>
|
|
64
|
+
|
|
65
|
+
---
|
|
66
|
+
|
|
67
|
+
Batched tensorsolve() for NumPy / PyTorch / JAX. ([numpy/numpy#28099](https://github.com/numpy/numpy/issues/28099))
|
|
68
|
+
|
|
69
|
+
## Installation
|
|
70
|
+
|
|
71
|
+
Install this via pip (or your favourite package manager):
|
|
72
|
+
|
|
73
|
+
```shell
|
|
74
|
+
pip install batch-tensorsolve
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
## Usage
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
import numpy as np
|
|
81
|
+
from numpy.testing import assert_allclose
|
|
82
|
+
|
|
83
|
+
from batch_tensorsolve import btensorsolve
|
|
84
|
+
|
|
85
|
+
a = np.random.randn(2, 2, 3, 6)
|
|
86
|
+
b = np.random.randn(2, 2, 3)
|
|
87
|
+
assert_allclose(np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b)), b)
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Advanced Usage
|
|
91
|
+
|
|
92
|
+
It is recommended to explicitly specify the batch axes, as the desired shape will be ambiguous if axes of size 1 are present.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
import numpy as np
|
|
96
|
+
|
|
97
|
+
from batch_tensorsolve import btensorsolve
|
|
98
|
+
|
|
99
|
+
a = np.random.randn(2, 1, 2, 2)
|
|
100
|
+
b = np.random.randn(2, 1, 2)
|
|
101
|
+
# 2 possibilities:
|
|
102
|
+
assert btensorsolve(a, b, num_batch_axes=1).shape == (2, 2) # 1st axis is batch
|
|
103
|
+
assert btensorsolve(a, b, num_batch_axes=2).shape == (2, 1, 2) # 1st and 2nd axes are batch
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Broadcasting-like behavior is also supported:
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
import numpy as np
|
|
110
|
+
from numpy.testing import assert_allclose
|
|
111
|
+
|
|
112
|
+
from batch_tensorsolve import btensorsolve
|
|
113
|
+
|
|
114
|
+
a = np.random.randn(1, 2, 3, 6) # -> (2, 2, 3, 6)
|
|
115
|
+
b = np.random.randn(2, 1, 1) # -> (2, 2, 3)
|
|
116
|
+
left = np.einsum("...ijk,...k->...ij", a, btensorsolve(a, b))
|
|
117
|
+
assert_allclose(left, np.broadcast_to(b, left.shape))
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
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.
|
|
121
|
+
|
|
122
|
+
## Contributors ✨
|
|
123
|
+
|
|
124
|
+
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
125
|
+
|
|
126
|
+
<!-- prettier-ignore-start -->
|
|
127
|
+
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
128
|
+
<!-- markdownlint-disable -->
|
|
129
|
+
<!-- markdownlint-enable -->
|
|
130
|
+
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
131
|
+
<!-- prettier-ignore-end -->
|
|
132
|
+
|
|
133
|
+
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
134
|
+
|
|
135
|
+
## Credits
|
|
136
|
+
|
|
137
|
+
[](https://github.com/copier-org/copier)
|
|
138
|
+
|
|
139
|
+
This package was created with
|
|
140
|
+
[Copier](https://copier.readthedocs.io/) and the
|
|
141
|
+
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
142
|
+
project template.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
setup.py
|
|
5
|
+
src/batch_tensorsolve/__init__.py
|
|
6
|
+
src/batch_tensorsolve/_main.py
|
|
7
|
+
src/batch_tensorsolve/py.typed
|
|
8
|
+
src/batch_tensorsolve.egg-info/PKG-INFO
|
|
9
|
+
src/batch_tensorsolve.egg-info/SOURCES.txt
|
|
10
|
+
src/batch_tensorsolve.egg-info/dependency_links.txt
|
|
11
|
+
src/batch_tensorsolve.egg-info/requires.txt
|
|
12
|
+
src/batch_tensorsolve.egg-info/top_level.txt
|
|
13
|
+
tests/test_main.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
batch_tensorsolve
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import pytest
|
|
3
|
+
from numpy.testing import assert_allclose
|
|
4
|
+
|
|
5
|
+
from batch_tensorsolve import (
|
|
6
|
+
AmbiguousBatchAxesWarning,
|
|
7
|
+
broadcast_without_repeating,
|
|
8
|
+
btensorsolve,
|
|
9
|
+
)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_btensorsolve():
|
|
13
|
+
rng = np.random.default_rng(0)
|
|
14
|
+
a = rng.normal(0, 1, (1, 1, 1, 2, 2, 3, 2, 6))
|
|
15
|
+
b = rng.normal(0, 1, (2, 1, 1, 1, 2, 3))
|
|
16
|
+
with pytest.warns(AmbiguousBatchAxesWarning):
|
|
17
|
+
sol = btensorsolve(a, b)
|
|
18
|
+
assert sol.shape == (2, 1, 1, 2, 6)
|
|
19
|
+
asol = np.einsum("...ijklm,...lm->...ijk", a, sol)
|
|
20
|
+
assert_allclose(asol, np.broadcast_to(b, asol.shape))
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def test_broadcast_without_repeating():
|
|
24
|
+
rng = np.random.default_rng(0)
|
|
25
|
+
a = rng.normal(0, 1, (3, 1))
|
|
26
|
+
b = rng.normal(0, 1, (2, 1, 4))
|
|
27
|
+
a, b = broadcast_without_repeating(a, b)
|
|
28
|
+
assert_allclose(a, np.broadcast_to(a, (1, 3, 1)))
|
|
29
|
+
assert_allclose(b, np.broadcast_to(b, (2, 1, 4)))
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def test_broadcast_without_repeating_error():
|
|
33
|
+
rng = np.random.default_rng(0)
|
|
34
|
+
a = rng.normal(0, 1, (2,))
|
|
35
|
+
b = rng.normal(0, 1, (3,))
|
|
36
|
+
with pytest.raises(ValueError):
|
|
37
|
+
broadcast_without_repeating(a, b)
|