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