feat: wip

Author: 34j

pyproject.toml

@@ -26,6 +26,8 @@ classifiers = [
 ]
 
 dependencies = [
+    "array-api-compat>=1.11.2",
+    "frozendict>=2.4.6",
 ]
 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 +36,11 @@ urls.repository = "https://github.com/34j/array-api-jit"
 
 [dependency-groups]
 dev = [
+  "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 +50,7 @@ docs = [
 ]
 
 [tool.ruff]
-line-length = 88
+line-length = 100
 lint.select = [
   "B",   # flake8-bugbear
   "D",   # flake8-docstrings

src/array_api_jit/main.py

@@ -1,3 +1,85 @@
-def add(n1: int, n2: int) -> int:
-    """Add the arguments."""
-    return n1 + n2
+# https://github.com/search?q=gumerov+translation+language%3APython&type=code&l=Python
+from collections.abc import Mapping, Sequence
+from functools import cache, wraps
+from types import ModuleType
+from typing import Any, Callable, ParamSpec, TypeVar
+
+from array_api_compat import array_namespace
+from frozendict import frozendict
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _default_decorator(
+    module: ModuleType, /, *args: Any, **kwargs: Any
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    if "numpy" in module.__name__:
+        import numba
+
+        return numba.jit(*args, **kwargs)
+    elif "torch" in module.__name__:
+        import torch
+
+        return torch.compile
+    else:
+        return getattr(module, "jit", lambda x: x)(*args, **kwargs)
+
+
+def jit(
+    decorator: Mapping[str, Callable[[Callable[P, T]], Callable[P, T]]] | None = None,
+    /,
+    *,
+    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
+    decorator_args : Mapping[str, Sequence[Any]] | None, optional
+        Additional positional arguments for the decorator for each array namespace, by default None
+    decorator_kwargs : Mapping[str, Mapping[str, Any]] | None, optional
+        Additional keyword arguments for 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.
+
+    """
+
+    def new_decorator(f: Callable[P, T]) -> Callable[P, T]:
+        decorator_args_ = frozendict(decorator_args or {})
+        decorator_kwargs_ = frozendict(decorator_kwargs or {})
+        decorator_ = decorator or {}
+
+        @cache
+        def jit_cached(xp: ModuleType) -> Callable[P, T]:
+            decorator_args__ = decorator_args_.get(xp, ())
+            decorator_kwargs__ = decorator_kwargs_.get(xp, {})
+            name = xp.__name__
+            if name in decorator_:
+                decorator_current = decorator_[name]
+            else:
+                decorator_current = _default_decorator(xp, *decorator_args_, **decorator_kwargs_)
+
+            return decorator_current(f, *decorator_args__, **decorator_kwargs__)
+
+        @wraps(f)
+        def inner(*args_inner: P.args, **kwargs_inner: P.kwargs) -> T:
+            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)
+            return f_jit(*args_inner, **kwargs_inner)
+
+        return inner
+
+    return new_decorator

tests/conftest.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+import pytest
+
+
+@pytest.fixture(scope="session", params=["numpy", "torch"])
+def xp(request: pytest.FixtureRequest) -> Any:
+    """Get the array namespace for the given backend."""
+    backend = request.param
+    if backend == "numpy":
+        import numpy as xp
+    elif backend == "torch":
+        import torch as xp
+    else:
+        raise ValueError(f"Unknown backend: {backend}")
+    return xp

tests/test_main.py

@@ -1,6 +1,26 @@
-from array_api_jit.main import add
+from typing import Any
 
+import array_api_compat.torch as torch
+from array_api_compat import array_namespace
 
-def test_add():
-    """Adding two number works as expected."""
-    assert add(1, 1) == 2
+from array_api_jit.main import jit
+
+
+def test_jit(xp: Any) -> None:
+    @jit(decorator_kwargs={"numpy": {"forceobj": True, "nopython": False}})  # type: ignore
+    def spherical_coordinates(
+        r: torch.Tensor, theta: torch.Tensor, phi: torch.Tensor
+    ) -> torch.Tensor:
+        xp = array_namespace(r, theta, phi)
+        rsin = r * xp.sin(theta)
+        x = rsin * xp.cos(phi)
+        y = rsin * xp.sin(phi)
+        z = r * xp.cos(theta)
+        return xp.stack((x, y, z), axis=-1)
+
+    r = xp.arange(100)
+    theta = xp.linspace(0, xp.pi, 100)
+    phi = xp.linspace(0, 2 * xp.pi, 100)
+    x = spherical_coordinates(r, theta, phi)  # type: ignore
+    print(x.shape)
+    assert x.shape == (100, 3)