array-api-jit 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

array_api_jit/_main.py

@@ -19,8 +19,8 @@ if importlib.util.find_spec("numba"):
     from numba.extending import overload
 
     @overload(array_namespace)
-    def _array_namespace_overload(*args: Any) -> Any:
-        def inner(*args: Any) -> Any:
+    def _array_namespace_overload(*_args: Any) -> Any:
+        def inner(*_args: Any) -> Any:
             return np
 
         return inner
@@ -41,6 +41,13 @@ STR_TO_IS_NAMESPACE = {
 }
 
 
+def _identity[**P, T](
+    f: Callable[P, T],
+    /,
+) -> Callable[P, T]:
+    return f
+
+
 def _default_decorator(
     module: ModuleType,
     /,
@@ -49,20 +56,19 @@ def _default_decorator(
         import jax
 
         return jax.jit
-    elif is_numpy_namespace(module) or is_cupy_namespace(module):
+    if is_numpy_namespace(module) or is_cupy_namespace(module):
         # import numba
 
         # return numba.jit()
         # The success rate of numba.jit is low
-        return lambda x: x
-    elif is_torch_namespace(module):
+        return _identity
+    if is_torch_namespace(module):
         import torch
 
         return torch.compile
-    elif is_dask_namespace(module):
-        return lambda x: x
-    else:
-        return getattr(module, "jit", lambda x: x)
+    if is_dask_namespace(module):
+        return _identity
+    return getattr(module, "jit", _identity)
 
 
 Decorator = Callable[[Callable[Pin, Tin]], Callable[Pin, Tin]]
@@ -120,7 +126,9 @@ def jit(
 
     """
 
-    def new_decorator(f: Callable[Pinner, Tinner]) -> Callable[Pinner, Tinner]:
+    def new_decorator(
+        f: Callable[Pinner, Tinner],
+    ) -> Callable[Pinner, Tinner]:
         decorator_args_ = decorator_args or {}
         decorator_kwargs_ = decorator_kwargs or {}
         decorator_ = decorator or {}
@@ -130,14 +138,10 @@ def jit(
             for name_, is_namespace in STR_TO_IS_NAMESPACE.items():
                 if is_namespace(xp):
                     name = name_
-            else:
-                name = xp.__name__.split(".")[0]
+            name = xp.__name__.split(".")[0]
             decorator_args__ = decorator_args_.get(name, ())
             decorator_kwargs__ = decorator_kwargs_.get(name, {})
-            if name in decorator_:
-                decorator_current = decorator_[name]
-            else:
-                decorator_current = _default_decorator(xp)
+            decorator_current = decorator_[name] if name in decorator_ else _default_decorator(xp)
             try:
                 return decorator_current(f, *decorator_args__, **decorator_kwargs__)
             except Exception as e:

{array_api_jit-1.0.0.dist-info → array_api_jit-1.1.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: array-api-jit
-Version: 1.0.0
+Version: 1.1.0
 Summary: JIT decorator supporting multiple array API compatible libraries
-Author-email: 34j <34j.95a2p@simplelogin.com>
-License: MIT
+Author: 34j
+License-Expression: MIT
 Project-URL: Bug Tracker, https://github.com/34j/array-api-jit/issues
 Project-URL: Changelog, https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md
 Project-URL: documentation, https://array-api-jit.readthedocs.io
@@ -12,12 +12,12 @@ 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 :: Only
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Software Development :: Libraries
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: array-api-compat>=1.11.2
@@ -32,19 +32,19 @@ Dynamic: license-file
   <a href="https://array-api-jit.readthedocs.io">
     <img src="https://img.shields.io/readthedocs/array-api-jit.svg?logo=read-the-docs&logoColor=fff&style=flat-square" alt="Documentation Status">
   </a>
+</p>
+<p align="center">
   <a href="https://codecov.io/gh/34j/array-api-jit">
     <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">
   </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 href="https://github.com/j178/prek">
+    <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/j178/prek/master/docs/assets/badge-v0.json" alt="prek">
   </a>
 </p>
 <p align="center">
@@ -78,6 +78,8 @@ pip install array-api-jit
 Simply decorate your function with `@jit()`:
 
 ```python
+from typing import Any
+
 from array_api_jit import jit
 
 
@@ -112,25 +114,3 @@ def sin_n_times(x: Any, n: int) -> Any:
         x = xp.sin(x)
     return x
 ```
-
-## 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.

array_api_jit-1.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+array_api_jit/__init__.py,sha256=6Cd5sMzZMC4jIJ97HL2eYIS8scDQonRkgK-uT221w_w,64
+array_api_jit/_main.py,sha256=5cjXZOyAl_m1NAE9YNUWD-wORnozj9GUl7gpEpq_UsI,5854
+array_api_jit/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+array_api_jit-1.1.0.dist-info/licenses/LICENSE,sha256=Otwb96xcK4QBqXXVdBKpLZVnc4mJuFTa8rsMUEpphMM,1061
+array_api_jit-1.1.0.dist-info/METADATA,sha256=Bks5iWQ_K91J7ktOSungqKarAt-VeVLHhiHwcGEKM88,4265
+array_api_jit-1.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+array_api_jit-1.1.0.dist-info/top_level.txt,sha256=90tmkew6kqpBt5VZbfvK2cvsOfCn_oFZfDf74F-s7Hg,14
+array_api_jit-1.1.0.dist-info/RECORD,,

{array_api_jit-1.0.0.dist-info → array_api_jit-1.1.0.dist-info}/WHEEL

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

array_api_jit-1.0.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-array_api_jit/__init__.py,sha256=6Cd5sMzZMC4jIJ97HL2eYIS8scDQonRkgK-uT221w_w,64
-array_api_jit/_main.py,sha256=RrE4K4y2wO4zukuRrddfBoGrojv8b5XopzMFqh8ftSU,5862
-array_api_jit/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-array_api_jit-1.0.0.dist-info/licenses/LICENSE,sha256=Otwb96xcK4QBqXXVdBKpLZVnc4mJuFTa8rsMUEpphMM,1061
-array_api_jit-1.0.0.dist-info/METADATA,sha256=9nup4sy6v_xHlno6M6HV4-1ZBCKiGxfINaruINoRnlI,5174
-array_api_jit-1.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-array_api_jit-1.0.0.dist-info/top_level.txt,sha256=90tmkew6kqpBt5VZbfvK2cvsOfCn_oFZfDf74F-s7Hg,14
-array_api_jit-1.0.0.dist-info/RECORD,,