chore: wip

Author: 34j

pyproject.toml

@@ -88,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" ]
 

src/array_api_jit/main.py

@@ -1,4 +1,5 @@
 # https://github.com/search?q=gumerov+translation+language%3APython&type=code&l=Python
+import importlib.util
 from collections.abc import Mapping, Sequence
 from functools import cache, wraps
 from types import ModuleType
@@ -7,23 +8,36 @@
 from array_api_compat import array_namespace
 from frozendict import frozendict
 
+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")
 
 
 def _default_decorator(
-    module: ModuleType, /, *args: Any, **kwargs: Any
+    module: ModuleType,
+    /,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     if "numpy" in module.__name__:
         import numba
 
-        return numba.jit(*args, **kwargs)
+        return numba.jit()
     elif "torch" in module.__name__:
         import torch
 
         return torch.compile
     else:
-        return getattr(module, "jit", lambda x: x)(*args, **kwargs)
+        return getattr(module, "jit", lambda x: x)
 
 
 def jit(
@@ -65,7 +79,7 @@ def jit_cached(xp: ModuleType) -> Callable[P, T]:
             if name in decorator_:
                 decorator_current = decorator_[name]
             else:
-                decorator_current = _default_decorator(xp, *decorator_args_, **decorator_kwargs_)
+                decorator_current = _default_decorator(xp)
 
             return decorator_current(f, *decorator_args__, **decorator_kwargs__)
 

tests/test_main.py

@@ -7,7 +7,7 @@
 
 
 def test_jit(xp: Any) -> None:
-    @jit(decorator_kwargs={"numpy": {"forceobj": True, "nopython": False}})  # type: ignore
+    @jit()  # type: ignore
     def spherical_coordinates(
         r: torch.Tensor, theta: torch.Tensor, phi: torch.Tensor
     ) -> torch.Tensor:
@@ -22,5 +22,4 @@ def spherical_coordinates(
     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)