jax.jit / torch.compile tests + fix torch detected mode (delegate detection to array-api-compat)

.github/workflows/ci.yml

@@ -36,6 +36,23 @@ jobs:
       - name: Fuzz
         run: make fuzz
 
+  compile:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install
+        run: |
+          python -m pip install -e ".[dev]"
+          python -m pip install jax
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+      - name: Import guard
+        run: python -c "import jax, torch; print('jax', jax.__version__, '| torch', torch.__version__)"
+      - name: Compilation tests
+        run: python -m pytest tests/backends/test_jax_torch_compilation.py -v -rs
+
   docs:
     runs-on: ubuntu-latest
     steps:

Makefile

@@ -1,4 +1,4 @@
-.PHONY: install fmt lint type test coverage check fuzz bench docs docs-serve notebook
+.PHONY: install fmt lint type test coverage check fuzz bench backends docs docs-serve notebook
 
 install:
 	python -m pip install -e ".[dev]"
@@ -33,6 +33,9 @@ bench:
 		--benchmark-only --benchmark-columns=min,median,ops \
 		--benchmark-group-by=func
 
+backends:
+	python -m pytest tests/backends/test_jax_torch_compilation.py -v -rs
+
 docs:
 	python -m zensical build --strict
 

README.md

@@ -33,7 +33,7 @@ print(psi.source)
 #             return 0.0
 #         return x * math.exp(-x)
 #     """
-#     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+#     xp = array_namespace(x)
 #     return xp.where(x < 0, 0.0, x * xp.exp(-x))
 ```
 

docs/development.md

@@ -44,6 +44,12 @@ numeric warnings (`overflow encountered`, `invalid value`,
 - **check** — the full `make check` gate on Python 3.12–3.14
   (Hypothesis derandomized).
 - **fuzz** — the five fixed fuzzer seeds.
+- **compile** — proves the generated source traces correctly under
+  `jax.jit` and `torch.compile` (jit/compile compatibility of the
+  generated code). Compiling detected-mode (unpinned) functions with
+  `torch.compile` additionally requires `array-api-compat>=1.15`
+  ("array_namespace can now be used under torch.compile"); pinned mode
+  (`namespace=`) has no such requirement.
 - **docs** — builds this site with `--strict`.
 - **pages** — runs the test suite and deploys this site to
   [GitHub Pages](https://grAItools.github.io/array-vectorize/) on every

docs/index.md

@@ -37,7 +37,7 @@ inspectable.
                 return 0.0
             return x * math.exp(-x)
         """
-        xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+        xp = array_namespace(x)
         return xp.where(x < 0, 0.0, x * xp.exp(-x))
     ```
 

pyproject.toml

@@ -10,7 +10,7 @@ readme = "README.md"
 requires-python = ">=3.12"
 license = { text = "BSD-3-Clause" }
 dependencies = [
-    "array-api-compat>=1.7",
+    "array-api-compat>=1.10",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -29,6 +29,10 @@ classifiers = [
 Source = "https://github.com/grAItools/array-vectorize"
 
 [project.optional-dependencies]
+backends = [
+    "jax>=0.4",
+    "torch>=2.2",
+]
 dev = [
     "pytest>=8.0",
     "pytest-benchmark>=4.0",
@@ -95,6 +99,9 @@ filterwarnings = [
     "ignore:overflow encountered:RuntimeWarning",
     "ignore:invalid value encountered:RuntimeWarning",
     "ignore:divide by zero encountered:RuntimeWarning",
+    # torch.compile (detected mode) traces array_api_compat's lru_cache-
+    # wrapped type dispatch on first compile; dynamo's informational notice
+    "ignore:Dynamo detected a call to a `functools.lru_cache`-wrapped function:UserWarning",
 ]
 markers = [
     "slow: performance tests that are timing-sensitive",

src/array_vectorize/codegen/emitter.py

@@ -142,25 +142,14 @@ def _gen_expr(node: Node, ns: str) -> ast.expr:
 
 
 def _namespace_line(ns: str, param_names: list[str]) -> ast.Assign:
-    # xp = array_namespace(*[a for a in (x, y) if hasattr(a, "__array_namespace__")])
-    # 'a' is comprehension-scoped, so it cannot collide with user names.
-    element = _load("a")
-    generator = ast.comprehension(
-        target=ast.Name(id="a", ctx=ast.Store()),
-        iter=ast.Tuple(elts=[_load(p) for p in param_names], ctx=ast.Load()),
-        ifs=[
-            ast.Call(
-                func=_load("hasattr"),
-                args=[_load("a"), ast.Constant(value="__array_namespace__")],
-                keywords=[],
-            )
-        ],
-        is_async=False,
-    )
-    filtered = ast.ListComp(elt=element, generators=[generator])
+    # xp = array_namespace(x, y)
+    # array_namespace itself ignores Python scalars/None (array-api-compat
+    # >= 1.10) and dispatches backend arrays that lack
+    # __array_namespace__ (torch.Tensor, CuPy), so every parameter is
+    # passed straight through; an all-scalar call still raises TypeError.
     call = ast.Call(
         func=_load("array_namespace"),
-        args=[ast.Starred(value=filtered, ctx=ast.Load())],
+        args=[_load(p) for p in param_names],
         keywords=[],
     )
     return ast.Assign(targets=[ast.Name(id=ns, ctx=ast.Store())], value=call)

src/array_vectorize/compat.py

@@ -5,12 +5,19 @@
 explicit namespace pins (``vectorize(namespace=...)``). They live here so
 none of their users duplicates them (and so api.py and fallback.py can
 share the check without an import cycle).
+
+Array detection is fully delegated to ``array_api_compat``: no homegrown
+``hasattr(value, "__array_namespace__")`` partial implementations of the
+standard's detection machinery, which miss backends whose arrays never
+exposed that attribute (torch.Tensor, CuPy).
 """
 
 from __future__ import annotations
 
 from typing import Any
 
+from array_api_compat import array_namespace
+
 __all__ = ["_check_namespace", "_is_array", "_py_scalar"]
 
 #: members every Array API namespace is expected to expose
@@ -35,7 +42,18 @@ def _check_namespace(namespace: Any) -> None:
 
 
 def _is_array(value: Any) -> bool:
-    return hasattr(value, "__array_namespace__")
+    """True if array-api-compat recognizes ``value`` as an array.
+
+    Fully delegated to ``array_namespace``: it ignores Python scalars/None,
+    dispatches torch tensors (which never exposed ``__array_namespace__``),
+    and duck-types any standard-compliant array — no homegrown hasattr
+    checks that miss backends.
+    """
+    try:
+        array_namespace(value)
+    except TypeError:
+        return False
+    return True
 
 
 def _py_scalar(value: Any) -> Any:
@@ -50,6 +68,8 @@ def _py_scalar(value: Any) -> Any:
     if item is not None:
         return item()
     dtype = getattr(value, "dtype", None)
+    # not backend detection: a conversion fallback for elements without
+    # .item() (array-api-strict); torch/numpy/jax elements all have .item()
     if hasattr(value, "__array_namespace__") and dtype is not None:
         name = str(dtype)
         if "bool" in name:

src/array_vectorize/frontend/extract.py

@@ -11,6 +11,7 @@
 from collections.abc import Callable
 from typing import Any
 
+from ..compat import _is_array
 from ..errors import _reject
 from .info import FunctionInfo, Param, ParamKind, _AstFunction
 from .lambda_id import _find_target
@@ -72,7 +73,7 @@ def _classify_closures(target: types.FunctionType, info: FunctionInfo) -> None:
             info.math_funcs[var] = attr
         elif isinstance(value, bool | int | float):
             info.closure_scalars[var] = value
-        elif hasattr(value, "__array_namespace__"):
+        elif _is_array(value):
             info.closure_arrays[var] = value
         elif isinstance(value, types.FunctionType):
             info.user_funcs[var] = value

src/array_vectorize/ir/ssa.py

@@ -13,10 +13,13 @@ def generated_name(name: str) -> str:
 
 
 #: Names the generated code cannot rename: ``array_namespace`` is imported
-#: by that name and ``hasattr`` is called by the namespace-detection
-#: comprehension. User variables with these base names are mangled with a
-#: trailing underscore. The namespace variable ``xp`` is NOT reserved: it is
-#: chosen dynamically to avoid user names (see _lower.lower_function).
+#: by that name. ``hasattr`` is kept reserved defensively: generated code no
+#: longer calls it (the namespace preamble passes arguments straight to
+#: array_namespace), but reserving it stays zero-risk and keeps older
+#: snapshots' name allocation stable. User variables with these base names
+#: are mangled with a trailing underscore. The namespace variable ``xp``
+#: is NOT reserved: it is chosen dynamically to avoid user names
+#: (see _lower.lower_function).
 RESERVED_NAMES = frozenset({"array_namespace", "hasattr"})
 
 

src/array_vectorize/verify.py

@@ -47,8 +47,11 @@ def verify_match(
         broadcast = xp.broadcast_arrays(*arrays)
         shape = broadcast[0].shape
         # constant-return functions produce a Python scalar or 0-d array;
-        # broadcast it to the input shape so element-wise comparison works
-        got_arr = got if hasattr(got, "shape") else xp.asarray(got)
+        # broadcast it to the input shape so element-wise comparison works.
+        # asarray on a plain Python scalar uses the backend's default dtype
+        # (torch: float32) — build it in float64 so precision survives to
+        # the comparison below
+        got_arr = got if hasattr(got, "shape") else xp.asarray(got, dtype=xp.float64)
         if got_arr.shape == ():
             got_arr = xp.broadcast_to(got_arr, shape)
         if tuple(got_arr.shape) != tuple(shape):
@@ -78,7 +81,10 @@ def verify_match(
         # compare values with their natural types (bools as 0/1) via float64;
         # coercing the expected side to the output dtype could hide miscompiles
         got_f = xp.astype(got_flat, xp.float64)
-        exp_f = xp.astype(xp.reshape(xp.asarray(expected_values), (-1,)), xp.float64)
+        # asarray on a list of Python scalars uses the backend's default
+        # dtype (torch: float32), which would lose precision BEFORE the
+        # astype below — build the expected side in float64 directly
+        exp_f = xp.reshape(xp.asarray(expected_values, dtype=xp.float64), (-1,))
         # exact equality first (covers matching infinities and identical values);
         # the tolerance term is gated on finiteness so inf-vs-finite mismatches
         # cannot slip through as inf <= inf