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

array_api_jit/__init__.py

@@ -1 +1,4 @@
-__version__ = "0.0.0"
+__version__ = "1.0.0"
+from ._main import jit
+
+__all__ = ["jit"]

array_api_jit/_main.py

@@ -0,0 +1,176 @@
+import importlib.util
+import warnings
+from collections.abc import Callable, Mapping, Sequence
+from functools import cache, wraps
+from types import ModuleType
+from typing import Any, ParamSpec, TypeVar
+
+from array_api_compat import (
+    array_namespace,
+    is_cupy_namespace,
+    is_dask_namespace,
+    is_jax_namespace,
+    is_numpy_namespace,
+    is_torch_namespace,
+)
+
+if importlib.util.find_spec("numba"):
+    import numpy as np
+    from numba.extending import overload
+
+    @overload(array_namespace)
+    def _array_namespace_overload(*args: Any) -> Any:
+        def inner(*args: Any) -> Any:
+            return np
+
+        return inner
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+Pin = ParamSpec("Pin")
+Tin = TypeVar("Tin")
+Pinner = ParamSpec("Pinner")
+Tinner = TypeVar("Tinner")
+STR_TO_IS_NAMESPACE = {
+    "numpy": is_numpy_namespace,
+    "jax": is_jax_namespace,
+    "cupy": is_cupy_namespace,
+    "torch": is_torch_namespace,
+    "dask": is_dask_namespace,
+}
+
+
+def _default_decorator(
+    module: ModuleType,
+    /,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    if is_jax_namespace(module):
+        import jax
+
+        return jax.jit
+    elif 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):
+        import torch
+
+        return torch.compile
+    elif is_dask_namespace(module):
+        return lambda x: x
+    else:
+        return getattr(module, "jit", lambda x: x)
+
+
+Decorator = Callable[[Callable[Pin, Tin]], Callable[Pin, Tin]]
+
+
+def jit(
+    decorator: Mapping[str, Decorator[..., Any]] | None = None,
+    /,
+    *,
+    fail_on_error: bool = False,
+    rerun_on_error: bool = False,
+    decorator_args: Mapping[str, Sequence[Any]] | None = None,
+    decorator_kwargs: Mapping[str, Mapping[str, Any]] | None = None,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    """
+    Just-in-time compilation decorator with multiple backends.
+
+    Parameters
+    ----------
+    decorator : Mapping[str, Callable[[Callable[P, T]], Callable[P, T]]] | None, optional
+        The JIT decorator to use for each array namespace, by default None
+    fail_on_error : bool, optional
+        If True, raise an error if the JIT decorator fails to apply.
+        If False, just warn and return the original function, by default False
+    rerun_on_error : bool, optional
+        If True, rerun the function without JIT if the function
+        with JIT applied fails, by default False
+    decorator_args : Mapping[str, Sequence[Any]] | None, optional
+        Additional positional arguments to be passed along with the function
+        to the decorator for each array namespace, by default None
+    decorator_kwargs : Mapping[str, Mapping[str, Any]] | None, optional
+        Additional keyword arguments to be passed along with the function
+        to the decorator for each array namespace, by default None
+
+    Returns
+    -------
+    Callable[[Callable[P, T]], Callable[P, T]]
+        The JIT decorator that can be applied to a function.
+
+    Example
+    -------
+    >>> from array_api_jit import jit
+    >>> from array_api_compat import array_namespace
+    >>> from typing import Any
+    >>> import numba
+    >>> @jit(
+    ...     {"numpy": numba.jit()},  # numba.jit is not used by default
+    ...     decorator_kwargs={"jax": {"static_argnames": ["n"]}},  # jax requires static_argnames
+    ... )
+    ... def sin_n_times(x: Any, n: int) -> Any:
+    ...     xp = array_namespace(x)
+    ...     for i in range(n):
+    ...         x = xp.sin(x)
+    ...     return x
+
+    """
+
+    def new_decorator(f: Callable[Pinner, Tinner]) -> Callable[Pinner, Tinner]:
+        decorator_args_ = decorator_args or {}
+        decorator_kwargs_ = decorator_kwargs or {}
+        decorator_ = decorator or {}
+
+        @cache
+        def jit_cached(xp: ModuleType) -> Callable[Pinner, Tinner]:
+            for name_, is_namespace in STR_TO_IS_NAMESPACE.items():
+                if is_namespace(xp):
+                    name = name_
+            else:
+                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)
+            try:
+                return decorator_current(f, *decorator_args__, **decorator_kwargs__)
+            except Exception as e:
+                if fail_on_error:
+                    raise RuntimeError(f"Failed to apply JIT decorator for {name}") from e
+                warnings.warn(
+                    f"Failed to apply JIT decorator for {name}: {e}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                return f
+
+        @wraps(f)
+        def inner(*args_inner: Pinner.args, **kwargs_inner: Pinner.kwargs) -> Tinner:
+            try:
+                xp = array_namespace(*args_inner)
+            except TypeError as e:
+                if e.args[0] == "Unrecognized array input":
+                    return f(*args_inner, **kwargs_inner)
+                raise
+            f_jit = jit_cached(xp)
+            try:
+                return f_jit(*args_inner, **kwargs_inner)
+            except Exception as e:
+                if rerun_on_error:
+                    warnings.warn(
+                        f"JIT failed for {xp.__name__}: {e}. Rerunning without JIT.",
+                        RuntimeWarning,
+                        stacklevel=2,
+                    )
+                    return f(*args_inner, **kwargs_inner)
+                raise RuntimeError(f"Failed to run JIT function for {xp.__name__}") from e
+
+        return inner
+
+    return new_decorator

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: array-api-jit
-Version: 0.0.0
+Version: 1.0.0
 Summary: JIT decorator supporting multiple array API compatible libraries
 Author-email: 34j <34j.95a2p@simplelogin.com>
 License: MIT
@@ -12,15 +12,15 @@ Classifier: Development Status :: 2 - Pre-Alpha
 Classifier: Intended Audience :: Developers
 Classifier: Natural Language :: English
 Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: array-api-compat>=1.11.2
 Dynamic: license-file
 
 # array API JIT
@@ -69,7 +69,49 @@ JIT decorator supporting multiple array API compatible libraries
 
 Install this via pip (or your favourite package manager):
 
-`pip install array-api-jit`
+```shell
+pip install array-api-jit
+```
+
+## Usage
+
+Simply decorate your function with `@jit()`:
+
+```python
+from array_api_jit import jit
+
+
+@jit()
+def my_function(x: Any) -> Any:
+    xp = array_namespace(x)
+    return xp.sin(x) + xp.cos(x)
+```
+
+## Advanced Usage
+
+You can specify the decorator, arguments, and keyword arguments for each library.
+
+```python
+from array_api_jit import jit
+from array_api_compat import array_namespace
+from typing import Any
+import numba
+
+
+@jit(
+    {"numpy": numba.jit()},  # numba.jit is not used by default because it may not succeed
+    decorator_kwargs={
+        "jax": {"static_argnames": ["n"]}
+    },  # jax requires for-loop variable to be "static_argnames"
+    # fail_on_error: bool = False, # do not raise an error if the decorator fails (Default)
+    # rerun_on_error: bool = True, # re-run the original function if the wrapped function fails (NOT Default)
+)
+def sin_n_times(x: Any, n: int) -> Any:
+    xp = array_namespace(x)
+    for i in range(n):
+        x = xp.sin(x)
+    return x
+```
 
 ## Contributors ✨
 

array_api_jit-1.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

array_api_jit/main.py

@@ -1,3 +0,0 @@
-def add(n1: int, n2: int) -> int:
-    """Add the arguments."""
-    return n1 + n2

array_api_jit-0.0.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-array_api_jit/__init__.py,sha256=ShXQBVjyiSOHxoQJS2BvNG395W4KZfqMxZWBAR0MZrE,22
-array_api_jit/main.py,sha256=j2h8YoaIyvMAosyxXlV3rlWEQ7x1yvnZWhdjVG4vXdU,82
-array_api_jit/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-array_api_jit-0.0.0.dist-info/licenses/LICENSE,sha256=Otwb96xcK4QBqXXVdBKpLZVnc4mJuFTa8rsMUEpphMM,1061
-array_api_jit-0.0.0.dist-info/METADATA,sha256=45CZdjO99TuZj8qkFY43C6yCaE23YMd9-iW7IKL2Uhw,4182
-array_api_jit-0.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-array_api_jit-0.0.0.dist-info/top_level.txt,sha256=90tmkew6kqpBt5VZbfvK2cvsOfCn_oFZfDf74F-s7Hg,14
-array_api_jit-0.0.0.dist-info/RECORD,,