array-api-jit 0.0.0__tar.gz → 1.0.0__tar.gz

{array_api_jit-0.0.0/src/array_api_jit.egg-info → array_api_jit-1.0.0}/PKG-INFO

@@ -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-0.0.0 → array_api_jit-1.0.0}/README.md

@@ -44,7 +44,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-0.0.0 → array_api_jit-1.0.0}/pyproject.toml

@@ -4,20 +4,19 @@ requires = [ "setuptools" ]
 
 [project]
 name = "array-api-jit"
-version = "0.0.0"
+version = "1.0.0"
 description = "JIT decorator supporting multiple array API compatible libraries"
 readme = "README.md"
 license = { text = "MIT" }
 authors = [
   { name = "34j", email = "34j.95a2p@simplelogin.com" },
 ]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 classifiers = [
   "Development Status :: 2 - Pre-Alpha",
   "Intended Audience :: Developers",
   "Natural Language :: English",
   "Operating System :: OS Independent",
-  "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
@@ -26,6 +25,7 @@ classifiers = [
 ]
 
 dependencies = [
+    "array-api-compat>=1.11.2",
 ]
 urls."Bug Tracker" = "https://github.com/34j/array-api-jit/issues"
 urls.Changelog = "https://github.com/34j/array-api-jit/blob/main/CHANGELOG.md"
@@ -34,8 +34,12 @@ urls.repository = "https://github.com/34j/array-api-jit"
 
 [dependency-groups]
 dev = [
+  "cm-time>=0.1.2",
+  "jax>=0.4.30",
+  "numba>=0.60.0",
   "pytest>=8,<9",
   "pytest-cov>=6,<7",
+  "torch>=2.7.1",
 ]
 docs = [
   "furo>=2023.5.20; python_version>='3.11'",
@@ -45,7 +49,8 @@ docs = [
 ]
 
 [tool.ruff]
-line-length = 88
+format.docstring-code-format = true
+line-length = 100
 lint.select = [
   "B",   # flake8-bugbear
   "D",   # flake8-docstrings
@@ -83,9 +88,6 @@ lint.isort.known-first-party = [ "array_api_jit", "tests" ]
 addopts = """\
     -v
     -Wdefault
-    --cov=array_api_jit
-    --cov-report=term
-    --cov-report=xml
     """
 pythonpath = [ "src" ]
 

array_api_jit-1.0.0/src/array_api_jit/__init__.py

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

array_api_jit-1.0.0/src/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 → array_api_jit-1.0.0/src/array_api_jit.egg-info}/PKG-INFO

@@ -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-0.0.0 → array_api_jit-1.0.0}/src/array_api_jit.egg-info/SOURCES.txt

@@ -3,10 +3,11 @@ README.md
 pyproject.toml
 setup.py
 src/array_api_jit/__init__.py
-src/array_api_jit/main.py
+src/array_api_jit/_main.py
 src/array_api_jit/py.typed
 src/array_api_jit.egg-info/PKG-INFO
 src/array_api_jit.egg-info/SOURCES.txt
 src/array_api_jit.egg-info/dependency_links.txt
+src/array_api_jit.egg-info/requires.txt
 src/array_api_jit.egg-info/top_level.txt
 tests/test_main.py

array_api_jit-1.0.0/src/array_api_jit.egg-info/requires.txt

@@ -0,0 +1 @@
+array-api-compat>=1.11.2

array_api_jit-1.0.0/tests/test_main.py

@@ -0,0 +1,133 @@
+from os import environ
+from sys import platform
+from typing import Any
+
+import numba
+import numpy as np
+import pytest
+from array_api_compat import array_namespace
+from cm_time import timer
+from numba import prange
+from numba.extending import overload
+
+from array_api_jit import jit
+
+IS_CI = environ.get("CI", "false").lower() in ("true", "1", "yes")
+
+
+@overload(np.stack)
+def stack(arrays, axis=0):
+    def inner(arrays, axis=0):
+        if axis == 0:
+            shape = (len(arrays),) + arrays[0].shape  # noqa: RUF005
+            stacked_array = np.empty(shape, dtype=arrays[0].dtype)
+            for j in prange(len(arrays)):
+                stacked_array[j] = arrays[j]
+        elif axis == -1:
+            shape = arrays[0].shape + (len(arrays),)  # noqa: RUF005
+            stacked_array = np.empty(shape, dtype=arrays[0].dtype)
+            for j in prange(len(arrays)):
+                stacked_array[..., j] = arrays[j]
+        return stacked_array
+
+    return inner
+
+
+def legendre(x: Any, n_end: int) -> Any:
+    """
+    Legendre polynomial of order 0 to n_end-1.
+
+    Parameters
+    ----------
+    x : Any
+        The points at which to evaluate the polynomial.
+    n_end : int
+        The number of orders to compute, starting from 0.
+
+    Returns
+    -------
+    Any
+        Array of shape (*x.shape, n_end),
+        where [..., i] is the value of the i-th order polynomial at x.
+
+    """
+    xp = array_namespace(x)
+    prevprev = xp.zeros_like(x)
+    prev = xp.ones_like(x)
+    result = [prevprev, prev]
+    for n in range(2, n_end):
+        prevprev, prev = prev, (((2 * n - 1) * x * prev - (n - 1) * prevprev) / n)
+        result.append(prev)
+    return xp.stack(result[:n_end], axis=-1)
+
+
+def legendre_assign(x: Any, n_end: int) -> Any:
+    """
+    Legendre polynomial of order 0 to n_end-1.
+
+    Parameters
+    ----------
+    x : Any
+        The points at which to evaluate the polynomial.
+    n_end : int
+        The number of orders to compute, starting from 0.
+
+    Returns
+    -------
+    Any
+        Array of shape (*x.shape, n_end),
+        where [..., i] is the value of the i-th order polynomial at x.
+
+    """
+    xp = array_namespace(x)
+    prevprev = xp.zeros_like(x)
+    prev = xp.ones_like(x)
+    result = xp.empty((*x.shape, n_end), dtype=x.dtype)
+    if n_end > 0:
+        result[..., 0] = prevprev
+    if n_end > 1:
+        result[..., 1] = prev
+    for n in range(2, n_end):
+        prevprev, prev = prev, (((2 * n - 1) * x * prev - (n - 1) * prevprev) / n)
+        result[..., n] = prev
+    return result
+
+
+legendre_jit = jit(
+    {"numpy": numba.jit(nogil=True, parallel=True)},
+    decorator_kwargs={"jax": {"static_argnames": ["n_end"]}},
+)(legendre)
+legendre_assign_jit = jit({"numpy": numba.jit(nogil=True, parallel=True)})(legendre_assign)
+
+
+@pytest.mark.skipif(
+    IS_CI and platform != "linux", reason="Compiler not available on GitHub Actions"
+)
+def test_jit(xp: Any) -> None:
+    t = {}
+    for name, func in [("nojit", legendre), ("jit", legendre_jit)] + (
+        [
+            ("assign-nojit", legendre_assign),
+            ("assign-jit", legendre_assign_jit),
+        ]
+        if "jax" not in xp.__name__
+        else []
+    ):
+        print(xp.__name__, name)
+        for i in range(20):
+            with timer() as timer_:
+                x = xp.arange(1000000, dtype=xp.float32) / 1000000
+                p = func(x, 10)
+                assert p.shape == (1000000, 10)
+            t[name] = timer_.elapsed
+            if i == 0:
+                print(f"First call: {timer_.elapsed:g}s")
+        print(f"Last call: {timer_.elapsed:g}s")
+    if "numpy" not in xp.__name__:
+        assert t["jit"] < t["nojit"], (
+            f"JIT time {t['jit']} should be less than non-JIT time {t['nojit']}"
+        )
+        assert t["assign-jit"] < t["assign-nojit"], (
+            f"JIT assign time {t['assign-jit']} should be "
+            f"less than non-JIT assign time {t['assign-nojit']}"
+        )

array_api_jit-0.0.0/src/array_api_jit/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.0"

array_api_jit-0.0.0/src/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/tests/test_main.py

@@ -1,6 +0,0 @@
-from array_api_jit.main import add
-
-
-def test_add():
-    """Adding two number works as expected."""
-    assert add(1, 1) == 2