chore: wip

Author: 34j

src/array_api_jit/main.py

@@ -6,7 +6,14 @@
 from types import ModuleType
 from typing import Any, Callable, ParamSpec, TypeVar
 
-from array_api_compat import array_namespace
+from array_api_compat import (
+    array_namespace,
+    is_cupy_namespace,
+    is_dask_namespace,
+    is_jax_namespace,
+    is_numpy_namespace,
+    is_torch_namespace,
+)
 from frozendict import frozendict
 
 if importlib.util.find_spec("numba"):
@@ -23,24 +30,35 @@ def inner(*args: Any) -> Any:
 
 P = ParamSpec("P")
 T = TypeVar("T")
+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 "jax" in module.__name__:
+    if is_jax_namespace(module):
         import jax
 
         return jax.jit
-    elif any(x in module.__name__ for x in ("numpy", "cupy")):
-        import numba
+    elif is_numpy_namespace(module) or is_cupy_namespace(module):
+        # import numba
 
-        return numba.jit()
-    elif "torch" in module.__name__:
+        # 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)
 
@@ -79,6 +97,22 @@ def jit(
     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 numpy as np
+    >>> @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[P, T]) -> Callable[P, T]:
@@ -88,9 +122,11 @@ def new_decorator(f: Callable[P, T]) -> Callable[P, T]:
 
         @cache
         def jit_cached(xp: ModuleType) -> Callable[P, T]:
-            name = xp.__name__
-            name = name.replace("array_api_compat.", "")
-            name = name.split(".")[0]
+            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_: