array-api-jit 1.0.0__tar.gz → 1.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.
- {array_api_jit-1.0.0/src/array_api_jit.egg-info → array_api_jit-1.1.0}/PKG-INFO +12 -32
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/README.md +6 -26
- array_api_jit-1.1.0/pyproject.toml +62 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit/_main.py +20 -16
- {array_api_jit-1.0.0 → array_api_jit-1.1.0/src/array_api_jit.egg-info}/PKG-INFO +12 -32
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit.egg-info/SOURCES.txt +0 -1
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/tests/test_main.py +6 -9
- array_api_jit-1.0.0/pyproject.toml +0 -162
- array_api_jit-1.0.0/setup.py +0 -9
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/LICENSE +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/setup.cfg +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit/__init__.py +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit/py.typed +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit.egg-info/dependency_links.txt +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit.egg-info/requires.txt +0 -0
- {array_api_jit-1.0.0 → array_api_jit-1.1.0}/src/array_api_jit.egg-info/top_level.txt +0 -0
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: array-api-jit
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.1.0
|
|
4
4
|
Summary: JIT decorator supporting multiple array API compatible libraries
|
|
5
|
-
Author
|
|
6
|
-
License: MIT
|
|
5
|
+
Author: 34j
|
|
6
|
+
License-Expression: MIT
|
|
7
7
|
Project-URL: Bug Tracker, https://github.com/34j/array-api-jit/issues
|
|
8
8
|
Project-URL: Changelog, https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md
|
|
9
9
|
Project-URL: documentation, https://array-api-jit.readthedocs.io
|
|
@@ -12,12 +12,12 @@ Classifier: Development Status :: 2 - Pre-Alpha
|
|
|
12
12
|
Classifier: Intended Audience :: Developers
|
|
13
13
|
Classifier: Natural Language :: English
|
|
14
14
|
Classifier: Operating System :: OS Independent
|
|
15
|
-
Classifier: Programming Language :: Python :: 3
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
17
16
|
Classifier: Programming Language :: Python :: 3.12
|
|
18
17
|
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
19
19
|
Classifier: Topic :: Software Development :: Libraries
|
|
20
|
-
Requires-Python: >=3.
|
|
20
|
+
Requires-Python: >=3.12
|
|
21
21
|
Description-Content-Type: text/markdown
|
|
22
22
|
License-File: LICENSE
|
|
23
23
|
Requires-Dist: array-api-compat>=1.11.2
|
|
@@ -32,19 +32,19 @@ Dynamic: license-file
|
|
|
32
32
|
<a href="https://array-api-jit.readthedocs.io">
|
|
33
33
|
<img src="https://img.shields.io/readthedocs/array-api-jit.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
34
34
|
</a>
|
|
35
|
+
</p>
|
|
36
|
+
<p align="center">
|
|
35
37
|
<a href="https://codecov.io/gh/34j/array-api-jit">
|
|
36
38
|
<img src="https://img.shields.io/codecov/c/github/34j/array-api-jit.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
37
39
|
</a>
|
|
38
|
-
</p>
|
|
39
|
-
<p align="center">
|
|
40
40
|
<a href="https://github.com/astral-sh/uv">
|
|
41
41
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
42
42
|
</a>
|
|
43
43
|
<a href="https://github.com/astral-sh/ruff">
|
|
44
44
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
45
45
|
</a>
|
|
46
|
-
<a href="https://github.com/
|
|
47
|
-
<img src="https://img.shields.io/
|
|
46
|
+
<a href="https://github.com/j178/prek">
|
|
47
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/j178/prek/master/docs/assets/badge-v0.json" alt="prek">
|
|
48
48
|
</a>
|
|
49
49
|
</p>
|
|
50
50
|
<p align="center">
|
|
@@ -78,6 +78,8 @@ pip install array-api-jit
|
|
|
78
78
|
Simply decorate your function with `@jit()`:
|
|
79
79
|
|
|
80
80
|
```python
|
|
81
|
+
from typing import Any
|
|
82
|
+
|
|
81
83
|
from array_api_jit import jit
|
|
82
84
|
|
|
83
85
|
|
|
@@ -112,25 +114,3 @@ def sin_n_times(x: Any, n: int) -> Any:
|
|
|
112
114
|
x = xp.sin(x)
|
|
113
115
|
return x
|
|
114
116
|
```
|
|
115
|
-
|
|
116
|
-
## Contributors ✨
|
|
117
|
-
|
|
118
|
-
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
119
|
-
|
|
120
|
-
<!-- prettier-ignore-start -->
|
|
121
|
-
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
122
|
-
<!-- markdownlint-disable -->
|
|
123
|
-
<!-- markdownlint-enable -->
|
|
124
|
-
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
125
|
-
<!-- prettier-ignore-end -->
|
|
126
|
-
|
|
127
|
-
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
128
|
-
|
|
129
|
-
## Credits
|
|
130
|
-
|
|
131
|
-
[](https://github.com/copier-org/copier)
|
|
132
|
-
|
|
133
|
-
This package was created with
|
|
134
|
-
[Copier](https://copier.readthedocs.io/) and the
|
|
135
|
-
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
136
|
-
project template.
|
|
@@ -7,19 +7,19 @@
|
|
|
7
7
|
<a href="https://array-api-jit.readthedocs.io">
|
|
8
8
|
<img src="https://img.shields.io/readthedocs/array-api-jit.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
9
9
|
</a>
|
|
10
|
+
</p>
|
|
11
|
+
<p align="center">
|
|
10
12
|
<a href="https://codecov.io/gh/34j/array-api-jit">
|
|
11
13
|
<img src="https://img.shields.io/codecov/c/github/34j/array-api-jit.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
12
14
|
</a>
|
|
13
|
-
</p>
|
|
14
|
-
<p align="center">
|
|
15
15
|
<a href="https://github.com/astral-sh/uv">
|
|
16
16
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
17
17
|
</a>
|
|
18
18
|
<a href="https://github.com/astral-sh/ruff">
|
|
19
19
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
20
20
|
</a>
|
|
21
|
-
<a href="https://github.com/
|
|
22
|
-
<img src="https://img.shields.io/
|
|
21
|
+
<a href="https://github.com/j178/prek">
|
|
22
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/j178/prek/master/docs/assets/badge-v0.json" alt="prek">
|
|
23
23
|
</a>
|
|
24
24
|
</p>
|
|
25
25
|
<p align="center">
|
|
@@ -53,6 +53,8 @@ pip install array-api-jit
|
|
|
53
53
|
Simply decorate your function with `@jit()`:
|
|
54
54
|
|
|
55
55
|
```python
|
|
56
|
+
from typing import Any
|
|
57
|
+
|
|
56
58
|
from array_api_jit import jit
|
|
57
59
|
|
|
58
60
|
|
|
@@ -87,25 +89,3 @@ def sin_n_times(x: Any, n: int) -> Any:
|
|
|
87
89
|
x = xp.sin(x)
|
|
88
90
|
return x
|
|
89
91
|
```
|
|
90
|
-
|
|
91
|
-
## Contributors ✨
|
|
92
|
-
|
|
93
|
-
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
94
|
-
|
|
95
|
-
<!-- prettier-ignore-start -->
|
|
96
|
-
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
97
|
-
<!-- markdownlint-disable -->
|
|
98
|
-
<!-- markdownlint-enable -->
|
|
99
|
-
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
100
|
-
<!-- prettier-ignore-end -->
|
|
101
|
-
|
|
102
|
-
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
103
|
-
|
|
104
|
-
## Credits
|
|
105
|
-
|
|
106
|
-
[](https://github.com/copier-org/copier)
|
|
107
|
-
|
|
108
|
-
This package was created with
|
|
109
|
-
[Copier](https://copier.readthedocs.io/) and the
|
|
110
|
-
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
111
|
-
project template.
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
build-backend = "setuptools.build_meta"
|
|
3
|
+
requires = [ "setuptools>=77.0.3" ]
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "array-api-jit"
|
|
7
|
+
version = "1.1.0"
|
|
8
|
+
description = "JIT decorator supporting multiple array API compatible libraries"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
authors = [
|
|
12
|
+
{ name = "34j" },
|
|
13
|
+
]
|
|
14
|
+
requires-python = ">=3.12"
|
|
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 :: Only",
|
|
21
|
+
"Programming Language :: Python :: 3.12",
|
|
22
|
+
"Programming Language :: Python :: 3.13",
|
|
23
|
+
"Programming Language :: Python :: 3.14",
|
|
24
|
+
"Topic :: Software Development :: Libraries",
|
|
25
|
+
]
|
|
26
|
+
dependencies = [
|
|
27
|
+
"array-api-compat>=1.11.2",
|
|
28
|
+
]
|
|
29
|
+
urls."Bug Tracker" = "https://github.com/34j/array-api-jit/issues"
|
|
30
|
+
urls.Changelog = "https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md"
|
|
31
|
+
urls.documentation = "https://array-api-jit.readthedocs.io"
|
|
32
|
+
urls.repository = "https://github.com/34j/array-api-jit"
|
|
33
|
+
|
|
34
|
+
[dependency-groups]
|
|
35
|
+
dev = [
|
|
36
|
+
"cm-time>=0.1.2",
|
|
37
|
+
"jax>=0.4.30",
|
|
38
|
+
"numba>=0.60",
|
|
39
|
+
"packaging-34j[test]>=1.1.8",
|
|
40
|
+
"torch>=2.7.1",
|
|
41
|
+
"ty>=0.0.99",
|
|
42
|
+
]
|
|
43
|
+
docs = [
|
|
44
|
+
"packaging-34j[docs]",
|
|
45
|
+
]
|
|
46
|
+
|
|
47
|
+
[tool.ruff]
|
|
48
|
+
extend = ".venv/ruff.toml"
|
|
49
|
+
lint.extend-per-file-ignores."/**/src/**/*" = [
|
|
50
|
+
"C901",
|
|
51
|
+
"PLC0415",
|
|
52
|
+
]
|
|
53
|
+
lint.extend-per-file-ignores."/**/tests/**/*" = [
|
|
54
|
+
"PLC0415",
|
|
55
|
+
"RUF005",
|
|
56
|
+
]
|
|
57
|
+
|
|
58
|
+
[tool.pyproject-fmt]
|
|
59
|
+
max_supported_python = "3.14"
|
|
60
|
+
|
|
61
|
+
[tool.ty]
|
|
62
|
+
terminal.error-on-warning = true
|
|
@@ -19,8 +19,8 @@ if importlib.util.find_spec("numba"):
|
|
|
19
19
|
from numba.extending import overload
|
|
20
20
|
|
|
21
21
|
@overload(array_namespace)
|
|
22
|
-
def _array_namespace_overload(*
|
|
23
|
-
def inner(*
|
|
22
|
+
def _array_namespace_overload(*_args: Any) -> Any:
|
|
23
|
+
def inner(*_args: Any) -> Any:
|
|
24
24
|
return np
|
|
25
25
|
|
|
26
26
|
return inner
|
|
@@ -41,6 +41,13 @@ STR_TO_IS_NAMESPACE = {
|
|
|
41
41
|
}
|
|
42
42
|
|
|
43
43
|
|
|
44
|
+
def _identity[**P, T](
|
|
45
|
+
f: Callable[P, T],
|
|
46
|
+
/,
|
|
47
|
+
) -> Callable[P, T]:
|
|
48
|
+
return f
|
|
49
|
+
|
|
50
|
+
|
|
44
51
|
def _default_decorator(
|
|
45
52
|
module: ModuleType,
|
|
46
53
|
/,
|
|
@@ -49,20 +56,19 @@ def _default_decorator(
|
|
|
49
56
|
import jax
|
|
50
57
|
|
|
51
58
|
return jax.jit
|
|
52
|
-
|
|
59
|
+
if is_numpy_namespace(module) or is_cupy_namespace(module):
|
|
53
60
|
# import numba
|
|
54
61
|
|
|
55
62
|
# return numba.jit()
|
|
56
63
|
# The success rate of numba.jit is low
|
|
57
|
-
return
|
|
58
|
-
|
|
64
|
+
return _identity
|
|
65
|
+
if is_torch_namespace(module):
|
|
59
66
|
import torch
|
|
60
67
|
|
|
61
68
|
return torch.compile
|
|
62
|
-
|
|
63
|
-
return
|
|
64
|
-
|
|
65
|
-
return getattr(module, "jit", lambda x: x)
|
|
69
|
+
if is_dask_namespace(module):
|
|
70
|
+
return _identity
|
|
71
|
+
return getattr(module, "jit", _identity)
|
|
66
72
|
|
|
67
73
|
|
|
68
74
|
Decorator = Callable[[Callable[Pin, Tin]], Callable[Pin, Tin]]
|
|
@@ -120,7 +126,9 @@ def jit(
|
|
|
120
126
|
|
|
121
127
|
"""
|
|
122
128
|
|
|
123
|
-
def new_decorator(
|
|
129
|
+
def new_decorator(
|
|
130
|
+
f: Callable[Pinner, Tinner],
|
|
131
|
+
) -> Callable[Pinner, Tinner]:
|
|
124
132
|
decorator_args_ = decorator_args or {}
|
|
125
133
|
decorator_kwargs_ = decorator_kwargs or {}
|
|
126
134
|
decorator_ = decorator or {}
|
|
@@ -130,14 +138,10 @@ def jit(
|
|
|
130
138
|
for name_, is_namespace in STR_TO_IS_NAMESPACE.items():
|
|
131
139
|
if is_namespace(xp):
|
|
132
140
|
name = name_
|
|
133
|
-
|
|
134
|
-
name = xp.__name__.split(".")[0]
|
|
141
|
+
name = xp.__name__.split(".")[0]
|
|
135
142
|
decorator_args__ = decorator_args_.get(name, ())
|
|
136
143
|
decorator_kwargs__ = decorator_kwargs_.get(name, {})
|
|
137
|
-
if name in decorator_
|
|
138
|
-
decorator_current = decorator_[name]
|
|
139
|
-
else:
|
|
140
|
-
decorator_current = _default_decorator(xp)
|
|
144
|
+
decorator_current = decorator_[name] if name in decorator_ else _default_decorator(xp)
|
|
141
145
|
try:
|
|
142
146
|
return decorator_current(f, *decorator_args__, **decorator_kwargs__)
|
|
143
147
|
except Exception as e:
|
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: array-api-jit
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.1.0
|
|
4
4
|
Summary: JIT decorator supporting multiple array API compatible libraries
|
|
5
|
-
Author
|
|
6
|
-
License: MIT
|
|
5
|
+
Author: 34j
|
|
6
|
+
License-Expression: MIT
|
|
7
7
|
Project-URL: Bug Tracker, https://github.com/34j/array-api-jit/issues
|
|
8
8
|
Project-URL: Changelog, https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md
|
|
9
9
|
Project-URL: documentation, https://array-api-jit.readthedocs.io
|
|
@@ -12,12 +12,12 @@ Classifier: Development Status :: 2 - Pre-Alpha
|
|
|
12
12
|
Classifier: Intended Audience :: Developers
|
|
13
13
|
Classifier: Natural Language :: English
|
|
14
14
|
Classifier: Operating System :: OS Independent
|
|
15
|
-
Classifier: Programming Language :: Python :: 3
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
17
16
|
Classifier: Programming Language :: Python :: 3.12
|
|
18
17
|
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
19
19
|
Classifier: Topic :: Software Development :: Libraries
|
|
20
|
-
Requires-Python: >=3.
|
|
20
|
+
Requires-Python: >=3.12
|
|
21
21
|
Description-Content-Type: text/markdown
|
|
22
22
|
License-File: LICENSE
|
|
23
23
|
Requires-Dist: array-api-compat>=1.11.2
|
|
@@ -32,19 +32,19 @@ Dynamic: license-file
|
|
|
32
32
|
<a href="https://array-api-jit.readthedocs.io">
|
|
33
33
|
<img src="https://img.shields.io/readthedocs/array-api-jit.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
|
|
34
34
|
</a>
|
|
35
|
+
</p>
|
|
36
|
+
<p align="center">
|
|
35
37
|
<a href="https://codecov.io/gh/34j/array-api-jit">
|
|
36
38
|
<img src="https://img.shields.io/codecov/c/github/34j/array-api-jit.svg?logo=codecov&logoColor=fff&style=flat-square" alt="Test coverage percentage">
|
|
37
39
|
</a>
|
|
38
|
-
</p>
|
|
39
|
-
<p align="center">
|
|
40
40
|
<a href="https://github.com/astral-sh/uv">
|
|
41
41
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json" alt="uv">
|
|
42
42
|
</a>
|
|
43
43
|
<a href="https://github.com/astral-sh/ruff">
|
|
44
44
|
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff">
|
|
45
45
|
</a>
|
|
46
|
-
<a href="https://github.com/
|
|
47
|
-
<img src="https://img.shields.io/
|
|
46
|
+
<a href="https://github.com/j178/prek">
|
|
47
|
+
<img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/j178/prek/master/docs/assets/badge-v0.json" alt="prek">
|
|
48
48
|
</a>
|
|
49
49
|
</p>
|
|
50
50
|
<p align="center">
|
|
@@ -78,6 +78,8 @@ pip install array-api-jit
|
|
|
78
78
|
Simply decorate your function with `@jit()`:
|
|
79
79
|
|
|
80
80
|
```python
|
|
81
|
+
from typing import Any
|
|
82
|
+
|
|
81
83
|
from array_api_jit import jit
|
|
82
84
|
|
|
83
85
|
|
|
@@ -112,25 +114,3 @@ def sin_n_times(x: Any, n: int) -> Any:
|
|
|
112
114
|
x = xp.sin(x)
|
|
113
115
|
return x
|
|
114
116
|
```
|
|
115
|
-
|
|
116
|
-
## Contributors ✨
|
|
117
|
-
|
|
118
|
-
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
|
|
119
|
-
|
|
120
|
-
<!-- prettier-ignore-start -->
|
|
121
|
-
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
|
|
122
|
-
<!-- markdownlint-disable -->
|
|
123
|
-
<!-- markdownlint-enable -->
|
|
124
|
-
<!-- ALL-CONTRIBUTORS-LIST:END -->
|
|
125
|
-
<!-- prettier-ignore-end -->
|
|
126
|
-
|
|
127
|
-
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
|
|
128
|
-
|
|
129
|
-
## Credits
|
|
130
|
-
|
|
131
|
-
[](https://github.com/copier-org/copier)
|
|
132
|
-
|
|
133
|
-
This package was created with
|
|
134
|
-
[Copier](https://copier.readthedocs.io/) and the
|
|
135
|
-
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
|
|
136
|
-
project template.
|
|
@@ -12,19 +12,19 @@ from numba.extending import overload
|
|
|
12
12
|
|
|
13
13
|
from array_api_jit import jit
|
|
14
14
|
|
|
15
|
-
IS_CI = environ.get("CI", "false").lower() in
|
|
15
|
+
IS_CI = environ.get("CI", "false").lower() in {"true", "1", "yes"}
|
|
16
16
|
|
|
17
17
|
|
|
18
18
|
@overload(np.stack)
|
|
19
19
|
def stack(arrays, axis=0):
|
|
20
20
|
def inner(arrays, axis=0):
|
|
21
21
|
if axis == 0:
|
|
22
|
-
shape = (len(arrays),) + arrays[0].shape
|
|
22
|
+
shape = (len(arrays),) + arrays[0].shape
|
|
23
23
|
stacked_array = np.empty(shape, dtype=arrays[0].dtype)
|
|
24
24
|
for j in prange(len(arrays)):
|
|
25
25
|
stacked_array[j] = arrays[j]
|
|
26
26
|
elif axis == -1:
|
|
27
|
-
shape = arrays[0].shape + (len(arrays),)
|
|
27
|
+
shape = arrays[0].shape + (len(arrays),)
|
|
28
28
|
stacked_array = np.empty(shape, dtype=arrays[0].dtype)
|
|
29
29
|
for j in prange(len(arrays)):
|
|
30
30
|
stacked_array[..., j] = arrays[j]
|
|
@@ -104,7 +104,7 @@ legendre_assign_jit = jit({"numpy": numba.jit(nogil=True, parallel=True)})(legen
|
|
|
104
104
|
IS_CI and platform != "linux", reason="Compiler not available on GitHub Actions"
|
|
105
105
|
)
|
|
106
106
|
def test_jit(xp: Any) -> None:
|
|
107
|
-
t = {}
|
|
107
|
+
t: dict[str, float] = {}
|
|
108
108
|
for name, func in [("nojit", legendre), ("jit", legendre_jit)] + (
|
|
109
109
|
[
|
|
110
110
|
("assign-nojit", legendre_assign),
|
|
@@ -113,16 +113,13 @@ def test_jit(xp: Any) -> None:
|
|
|
113
113
|
if "jax" not in xp.__name__
|
|
114
114
|
else []
|
|
115
115
|
):
|
|
116
|
-
|
|
117
|
-
for i in range(20):
|
|
116
|
+
for _ in range(20):
|
|
118
117
|
with timer() as timer_:
|
|
119
118
|
x = xp.arange(1000000, dtype=xp.float32) / 1000000
|
|
120
119
|
p = func(x, 10)
|
|
121
120
|
assert p.shape == (1000000, 10)
|
|
121
|
+
assert timer_.elapsed is not None
|
|
122
122
|
t[name] = timer_.elapsed
|
|
123
|
-
if i == 0:
|
|
124
|
-
print(f"First call: {timer_.elapsed:g}s")
|
|
125
|
-
print(f"Last call: {timer_.elapsed:g}s")
|
|
126
123
|
if "numpy" not in xp.__name__:
|
|
127
124
|
assert t["jit"] < t["nojit"], (
|
|
128
125
|
f"JIT time {t['jit']} should be less than non-JIT time {t['nojit']}"
|
|
@@ -1,162 +0,0 @@
|
|
|
1
|
-
[build-system]
|
|
2
|
-
build-backend = "setuptools.build_meta"
|
|
3
|
-
requires = [ "setuptools" ]
|
|
4
|
-
|
|
5
|
-
[project]
|
|
6
|
-
name = "array-api-jit"
|
|
7
|
-
version = "1.0.0"
|
|
8
|
-
description = "JIT decorator supporting multiple array API compatible libraries"
|
|
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
|
-
]
|
|
30
|
-
urls."Bug Tracker" = "https://github.com/34j/array-api-jit/issues"
|
|
31
|
-
urls.Changelog = "https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md"
|
|
32
|
-
urls.documentation = "https://array-api-jit.readthedocs.io"
|
|
33
|
-
urls.repository = "https://github.com/34j/array-api-jit"
|
|
34
|
-
|
|
35
|
-
[dependency-groups]
|
|
36
|
-
dev = [
|
|
37
|
-
"cm-time>=0.1.2",
|
|
38
|
-
"jax>=0.4.30",
|
|
39
|
-
"numba>=0.60.0",
|
|
40
|
-
"pytest>=8,<9",
|
|
41
|
-
"pytest-cov>=6,<7",
|
|
42
|
-
"torch>=2.7.1",
|
|
43
|
-
]
|
|
44
|
-
docs = [
|
|
45
|
-
"furo>=2023.5.20; python_version>='3.11'",
|
|
46
|
-
"myst-parser>=0.16; python_version>='3.11'",
|
|
47
|
-
"sphinx>=4; python_version>='3.11'",
|
|
48
|
-
"sphinx-autobuild>=2024,<2025; python_version>='3.11'",
|
|
49
|
-
]
|
|
50
|
-
|
|
51
|
-
[tool.ruff]
|
|
52
|
-
format.docstring-code-format = true
|
|
53
|
-
line-length = 100
|
|
54
|
-
lint.select = [
|
|
55
|
-
"B", # flake8-bugbear
|
|
56
|
-
"D", # flake8-docstrings
|
|
57
|
-
"C4", # flake8-comprehensions
|
|
58
|
-
"S", # flake8-bandit
|
|
59
|
-
"F", # pyflake
|
|
60
|
-
"E", # pycodestyle
|
|
61
|
-
"W", # pycodestyle
|
|
62
|
-
"UP", # pyupgrade
|
|
63
|
-
"I", # isort
|
|
64
|
-
"RUF", # ruff specific
|
|
65
|
-
]
|
|
66
|
-
lint.ignore = [
|
|
67
|
-
"D203", # 1 blank line required before class docstring
|
|
68
|
-
"D212", # Multi-line docstring summary should start at the first line
|
|
69
|
-
"D100", # Missing docstring in public module
|
|
70
|
-
"D104", # Missing docstring in public package
|
|
71
|
-
"D107", # Missing docstring in `__init__`
|
|
72
|
-
"D401", # First line of docstring should be in imperative mood
|
|
73
|
-
]
|
|
74
|
-
lint.per-file-ignores."conftest.py" = [ "D100" ]
|
|
75
|
-
lint.per-file-ignores."docs/conf.py" = [ "D100" ]
|
|
76
|
-
lint.per-file-ignores."setup.py" = [ "D100" ]
|
|
77
|
-
lint.per-file-ignores."tests/**/*" = [
|
|
78
|
-
"D100",
|
|
79
|
-
"D101",
|
|
80
|
-
"D102",
|
|
81
|
-
"D103",
|
|
82
|
-
"D104",
|
|
83
|
-
"S101",
|
|
84
|
-
]
|
|
85
|
-
lint.isort.known-first-party = [ "array_api_jit", "tests" ]
|
|
86
|
-
|
|
87
|
-
[tool.pytest.ini_options]
|
|
88
|
-
addopts = """\
|
|
89
|
-
-v
|
|
90
|
-
-Wdefault
|
|
91
|
-
"""
|
|
92
|
-
pythonpath = [ "src" ]
|
|
93
|
-
|
|
94
|
-
[tool.coverage.run]
|
|
95
|
-
branch = true
|
|
96
|
-
|
|
97
|
-
[tool.coverage.report]
|
|
98
|
-
exclude_lines = [
|
|
99
|
-
"pragma: no cover",
|
|
100
|
-
"@overload",
|
|
101
|
-
"if TYPE_CHECKING",
|
|
102
|
-
"raise NotImplementedError",
|
|
103
|
-
'if __name__ == "__main__":',
|
|
104
|
-
]
|
|
105
|
-
|
|
106
|
-
[tool.mypy]
|
|
107
|
-
check_untyped_defs = true
|
|
108
|
-
disallow_any_generics = true
|
|
109
|
-
disallow_incomplete_defs = true
|
|
110
|
-
disallow_untyped_defs = true
|
|
111
|
-
mypy_path = "src/"
|
|
112
|
-
no_implicit_optional = true
|
|
113
|
-
show_error_codes = true
|
|
114
|
-
warn_unreachable = true
|
|
115
|
-
warn_unused_ignores = true
|
|
116
|
-
exclude = [
|
|
117
|
-
'docs/.*',
|
|
118
|
-
'setup.py',
|
|
119
|
-
]
|
|
120
|
-
|
|
121
|
-
[[tool.mypy.overrides]]
|
|
122
|
-
module = "tests.*"
|
|
123
|
-
allow_untyped_defs = true
|
|
124
|
-
|
|
125
|
-
[[tool.mypy.overrides]]
|
|
126
|
-
module = "docs.*"
|
|
127
|
-
ignore_errors = true
|
|
128
|
-
|
|
129
|
-
[tool.semantic_release]
|
|
130
|
-
version_toml = [ "pyproject.toml:project.version" ]
|
|
131
|
-
version_variables = [
|
|
132
|
-
"src/array_api_jit/__init__.py:__version__",
|
|
133
|
-
"docs/conf.py:release",
|
|
134
|
-
]
|
|
135
|
-
build_command = """
|
|
136
|
-
pip install uv
|
|
137
|
-
uv lock
|
|
138
|
-
git add uv.lock
|
|
139
|
-
uv build
|
|
140
|
-
"""
|
|
141
|
-
|
|
142
|
-
[tool.semantic_release.changelog]
|
|
143
|
-
exclude_commit_patterns = [
|
|
144
|
-
'''chore(?:\([^)]*?\))?: .+''',
|
|
145
|
-
'''ci(?:\([^)]*?\))?: .+''',
|
|
146
|
-
'''refactor(?:\([^)]*?\))?: .+''',
|
|
147
|
-
'''style(?:\([^)]*?\))?: .+''',
|
|
148
|
-
'''test(?:\([^)]*?\))?: .+''',
|
|
149
|
-
'''build\((?!deps\): .+)''',
|
|
150
|
-
'''Merged? .*''',
|
|
151
|
-
'''Initial [Cc]ommit.*''', # codespell:ignore
|
|
152
|
-
]
|
|
153
|
-
|
|
154
|
-
[tool.semantic_release.changelog.environment]
|
|
155
|
-
keep_trailing_newline = true
|
|
156
|
-
|
|
157
|
-
[tool.semantic_release.branches.main]
|
|
158
|
-
match = "main"
|
|
159
|
-
|
|
160
|
-
[tool.semantic_release.branches.noop]
|
|
161
|
-
match = "(?!main$)"
|
|
162
|
-
prerelease = true
|
array_api_jit-1.0.0/setup.py
DELETED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|