From c91597974bb5b15464cab7107e139b5cc3cecb55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= <enriqueg@cscs.ch>
Date: Mon, 5 Oct 2026 21:50:00 +0200
Subject: [PATCH 1/4] codegen: pass args directly to array_namespace (fix torch
 detected mode)

---
 README.md                                  |  2 +-
 docs/index.md                              |  2 +-
 pyproject.toml                             |  2 +-
 src/array_vectorize/codegen/emitter.py     | 23 ++++++----------------
 src/array_vectorize/ir/ssa.py              | 11 +++++++----
 tests/features/test_namespace_pin.py       |  5 ++++-
 tests/golden/cases/abs_round.py            |  2 +-
 tests/golden/cases/add.py                  |  2 +-
 tests/golden/cases/arith_ops.py            |  2 +-
 tests/golden/cases/bitwise_ops.py          |  2 +-
 tests/golden/cases/boolop_logical.py       |  2 +-
 tests/golden/cases/boolop_numeric.py       |  2 +-
 tests/golden/cases/both_branches_assign.py |  2 +-
 tests/golden/cases/casts.py                |  2 +-
 tests/golden/cases/clamp.py                |  2 +-
 tests/golden/cases/closure_array.py        |  2 +-
 tests/golden/cases/closure_scalar.py       |  2 +-
 tests/golden/cases/compare_chain.py        |  2 +-
 tests/golden/cases/compare_ops.py          |  2 +-
 tests/golden/cases/complex_expr.py         |  2 +-
 tests/golden/cases/constants.py            |  2 +-
 tests/golden/cases/cse_opportunity.py      |  2 +-
 tests/golden/cases/floor_mod_ops.py        |  2 +-
 tests/golden/cases/from_import.py          |  2 +-
 tests/golden/cases/helper_outer.py         |  2 +-
 tests/golden/cases/kwonly_defaults.py      |  2 +-
 tests/golden/cases/loop_accumulate.py      |  2 +-
 tests/golden/cases/loop_helper_caller.py   |  2 +-
 tests/golden/cases/loop_start_stop_step.py |  2 +-
 tests/golden/cases/loop_two_carried.py     |  2 +-
 tests/golden/cases/loop_with_branch.py     |  2 +-
 tests/golden/cases/math_calls.py           |  2 +-
 tests/golden/cases/math_floor_ceil.py      |  2 +-
 tests/golden/cases/math_two_arg.py         |  2 +-
 tests/golden/cases/merge_with_prior.py     |  2 +-
 tests/golden/cases/minmax.py               |  2 +-
 tests/golden/cases/nested_early_returns.py |  2 +-
 tests/golden/cases/nested_loops.py         |  2 +-
 tests/golden/cases/piecewise.py            |  2 +-
 tests/golden/cases/pow_op.py               |  2 +-
 tests/golden/cases/psi.py                  |  2 +-
 tests/golden/cases/relu.py                 |  2 +-
 tests/golden/cases/relu_with_else.py       |  2 +-
 tests/golden/cases/reserved_param.py       |  2 +-
 tests/golden/cases/ssa_rebind.py           |  2 +-
 tests/golden/cases/ternary.py              |  2 +-
 tests/golden/cases/unary_ops.py            |  2 +-
 tests/unit/test_codegen.py                 |  4 ++--
 48 files changed, 63 insertions(+), 68 deletions(-)

diff --git a/README.md b/README.md
index f2bc2e1..b82cba6 100644
--- a/README.md
+++ b/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))
 ```
 
diff --git a/docs/index.md b/docs/index.md
index b291c7a..01f837b 100644
--- a/docs/index.md
+++ b/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))
     ```
 
diff --git a/pyproject.toml b/pyproject.toml
index 3d97c31..42933f6 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py
index 3ef51ed..5166016 100644
--- a/src/array_vectorize/codegen/emitter.py
+++ b/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)
diff --git a/src/array_vectorize/ir/ssa.py b/src/array_vectorize/ir/ssa.py
index 06660f7..956b651 100644
--- a/src/array_vectorize/ir/ssa.py
+++ b/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"})
 
 
diff --git a/tests/features/test_namespace_pin.py b/tests/features/test_namespace_pin.py
index f1af20f..e9e1537 100644
--- a/tests/features/test_namespace_pin.py
+++ b/tests/features/test_namespace_pin.py
@@ -91,7 +91,10 @@ def test_all_scalar_minmax_with_pin() -> None:
 def test_all_scalar_call_without_pin_raises() -> None:
     fn = make_fn("    return math.sqrt(x * x)", defaults="x")
     vec = vectorize(fn)
-    with pytest.raises(TypeError, match="array_namespace requires at least one non-scalar"):
+    # message wording varies across array-api-compat versions: 1.10 says
+    # "Unrecognized array input", newer releases "requires at least one
+    # non-scalar array input" — the contract is the TypeError itself
+    with pytest.raises(TypeError, match=r"at least one non-scalar|Unrecognized array input"):
         vec(3.0)
 
 
diff --git a/tests/golden/cases/abs_round.py b/tests/golden/cases/abs_round.py
index d9a67d5..f298459 100644
--- a/tests/golden/cases/abs_round.py
+++ b/tests/golden/cases/abs_round.py
@@ -5,6 +5,6 @@ def abs_round_vec(x):
     def abs_round(x):
         return abs(x) + round(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     t_1 = xp.asarray(x)
     return xp.abs(t_1) + xp.round(t_1)
diff --git a/tests/golden/cases/add.py b/tests/golden/cases/add.py
index 80bffb9..eb32789 100644
--- a/tests/golden/cases/add.py
+++ b/tests/golden/cases/add.py
@@ -5,5 +5,5 @@ def add_vec(x, y):
     def add(x, y):
         return x + y
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return x + y
diff --git a/tests/golden/cases/arith_ops.py b/tests/golden/cases/arith_ops.py
index 66e39a2..d315119 100644
--- a/tests/golden/cases/arith_ops.py
+++ b/tests/golden/cases/arith_ops.py
@@ -5,5 +5,5 @@ def arith_ops_vec(x, y):
     def arith_ops(x, y):
         return (x + y - x * y) / (x * x + y * y)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return (x + y - x * y) / (x * x + y * y)
diff --git a/tests/golden/cases/bitwise_ops.py b/tests/golden/cases/bitwise_ops.py
index 0406f21..1c9c702 100644
--- a/tests/golden/cases/bitwise_ops.py
+++ b/tests/golden/cases/bitwise_ops.py
@@ -5,5 +5,5 @@ def bitwise_ops_vec(x, y):
     def bitwise_ops(x, y):
         return (x & y) | (x ^ 3) << 1 >> 2
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return x & y | (x ^ 3) << 1 >> 2
diff --git a/tests/golden/cases/boolop_logical.py b/tests/golden/cases/boolop_logical.py
index b44f6f7..4f95aff 100644
--- a/tests/golden/cases/boolop_logical.py
+++ b/tests/golden/cases/boolop_logical.py
@@ -5,5 +5,5 @@ def boolop_logical_vec(x, y):
     def boolop_logical(x, y):
         return ((x > 0) and (y > 0)) or not (x < 0)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return xp.logical_or(xp.logical_and(x > 0, y > 0), xp.logical_not(x < 0))
diff --git a/tests/golden/cases/boolop_numeric.py b/tests/golden/cases/boolop_numeric.py
index 336077b..a3f81d0 100644
--- a/tests/golden/cases/boolop_numeric.py
+++ b/tests/golden/cases/boolop_numeric.py
@@ -5,5 +5,5 @@ def boolop_numeric_vec(x, y):
     def boolop_numeric(x, y):
         return x and y
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return xp.where(x != 0, y, x)
diff --git a/tests/golden/cases/both_branches_assign.py b/tests/golden/cases/both_branches_assign.py
index 2a3ed20..e0c4a59 100644
--- a/tests/golden/cases/both_branches_assign.py
+++ b/tests/golden/cases/both_branches_assign.py
@@ -11,7 +11,7 @@ def both_branches_assign(x):
             z = 2.0
         return y + z
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     y_1 = x
     z_1 = 1.0
     y_2 = -x
diff --git a/tests/golden/cases/casts.py b/tests/golden/cases/casts.py
index 1befbfe..4ee85ca 100644
--- a/tests/golden/cases/casts.py
+++ b/tests/golden/cases/casts.py
@@ -5,7 +5,7 @@ def casts_vec(x):
     def casts(x):
         return int(x) + float(x) + bool(x) + math.trunc(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.int64)
     return _vec_arith(xp, 1, t_2 + xp.astype(t_1, xp.float64), xp.astype(t_1, xp.bool)) + t_2
diff --git a/tests/golden/cases/clamp.py b/tests/golden/cases/clamp.py
index 904df34..39406fc 100644
--- a/tests/golden/cases/clamp.py
+++ b/tests/golden/cases/clamp.py
@@ -9,5 +9,5 @@ def clamp(x, lo=0.0, hi=1.0):
             return hi
         return x
     """
-    xp = array_namespace(*[a for a in (x, lo, hi) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, lo, hi)
     return xp.where(x < lo, lo, xp.where(x > hi, hi, x))
diff --git a/tests/golden/cases/closure_array.py b/tests/golden/cases/closure_array.py
index b657fd0..61b5857 100644
--- a/tests/golden/cases/closure_array.py
+++ b/tests/golden/cases/closure_array.py
@@ -5,5 +5,5 @@ def closure_array_vec(x, *, ARR=None):
     def closure_array(x):
         return x + ARR
     """
-    xp = array_namespace(*[a for a in (x, ARR) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, ARR)
     return x + ARR
diff --git a/tests/golden/cases/closure_scalar.py b/tests/golden/cases/closure_scalar.py
index 321db23..317b251 100644
--- a/tests/golden/cases/closure_scalar.py
+++ b/tests/golden/cases/closure_scalar.py
@@ -5,5 +5,5 @@ def closure_scalar_vec(x):
     def closure_scalar(x):
         return x * SCALE
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return x * 3.0
diff --git a/tests/golden/cases/compare_chain.py b/tests/golden/cases/compare_chain.py
index 702fdc4..9f8c728 100644
--- a/tests/golden/cases/compare_chain.py
+++ b/tests/golden/cases/compare_chain.py
@@ -5,5 +5,5 @@ def compare_chain_vec(x):
     def compare_chain(x):
         return 0 < x < 1
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.logical_and(0 < x, x < 1)
diff --git a/tests/golden/cases/compare_ops.py b/tests/golden/cases/compare_ops.py
index b5f643d..c874f9d 100644
--- a/tests/golden/cases/compare_ops.py
+++ b/tests/golden/cases/compare_ops.py
@@ -5,5 +5,5 @@ def compare_ops_vec(x, y):
     def compare_ops(x, y):
         return (x == y) + (x != y) + (x < y) + (x <= y) + (x > y) + (x >= y)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     return _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, x == y, x != y), x < y), x <= y), x > y), x >= y)
diff --git a/tests/golden/cases/complex_expr.py b/tests/golden/cases/complex_expr.py
index da9c8c3..91613ab 100644
--- a/tests/golden/cases/complex_expr.py
+++ b/tests/golden/cases/complex_expr.py
@@ -6,6 +6,6 @@ def complex_expr(x, y):
         t = math.exp(-(x * x + y * y))
         return t / (1.0 + t)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     t = xp.exp(xp.astype(xp.asarray(-(x * x + y * y)), xp.float64))
     return t / (1.0 + t)
diff --git a/tests/golden/cases/constants.py b/tests/golden/cases/constants.py
index f6c4e49..6d24dbf 100644
--- a/tests/golden/cases/constants.py
+++ b/tests/golden/cases/constants.py
@@ -5,5 +5,5 @@ def constants_vec(x):
     def constants(x):
         return x * math.pi + math.e + math.inf - math.nan
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return x * 3.141592653589793 + 2.718281828459045 + xp.inf - xp.nan
diff --git a/tests/golden/cases/cse_opportunity.py b/tests/golden/cases/cse_opportunity.py
index acffa8c..1fdb76b 100644
--- a/tests/golden/cases/cse_opportunity.py
+++ b/tests/golden/cases/cse_opportunity.py
@@ -5,7 +5,7 @@ def cse_opportunity_vec(x):
     def cse_opportunity(x):
         return math.sqrt(x) * math.sqrt(x) + math.sqrt(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.float64)
     t_3 = xp.sqrt(t_2)
diff --git a/tests/golden/cases/floor_mod_ops.py b/tests/golden/cases/floor_mod_ops.py
index c1a0678..66789f5 100644
--- a/tests/golden/cases/floor_mod_ops.py
+++ b/tests/golden/cases/floor_mod_ops.py
@@ -5,5 +5,5 @@ def floor_mod_ops_vec(x):
     def floor_mod_ops(x):
         return (x // 2) + (x % 3)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.floor_divide(x, 2) + xp.remainder(x, 3)
diff --git a/tests/golden/cases/from_import.py b/tests/golden/cases/from_import.py
index 000b4d7..39870f2 100644
--- a/tests/golden/cases/from_import.py
+++ b/tests/golden/cases/from_import.py
@@ -5,5 +5,5 @@ def from_import_vec(x):
     def from_import(x):
         return fexp(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.exp(xp.astype(xp.asarray(x), xp.float64))
diff --git a/tests/golden/cases/helper_outer.py b/tests/golden/cases/helper_outer.py
index 078d866..342406b 100644
--- a/tests/golden/cases/helper_outer.py
+++ b/tests/golden/cases/helper_outer.py
@@ -5,5 +5,5 @@ def helper_outer_vec(x):
     def helper_outer(x):
         return helper_inner(x) + helper_inner(x * 2.0)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return helper_inner_vec(x) + helper_inner_vec(x * 2.0)
diff --git a/tests/golden/cases/kwonly_defaults.py b/tests/golden/cases/kwonly_defaults.py
index f7c76dd..a4f7914 100644
--- a/tests/golden/cases/kwonly_defaults.py
+++ b/tests/golden/cases/kwonly_defaults.py
@@ -5,5 +5,5 @@ def kwonly_defaults_vec(x, *, scale=2.0, bias=0.5):
     def kwonly_defaults(x, *, scale=2.0, bias=0.5):
         return x * scale + bias
     """
-    xp = array_namespace(*[a for a in (x, scale, bias) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, scale, bias)
     return x * scale + bias
diff --git a/tests/golden/cases/loop_accumulate.py b/tests/golden/cases/loop_accumulate.py
index 13f886a..5abee99 100644
--- a/tests/golden/cases/loop_accumulate.py
+++ b/tests/golden/cases/loop_accumulate.py
@@ -8,7 +8,7 @@ def loop_accumulate(x):
             s = s + x
         return s
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     s = 0.0
     s_1 = s
     for i in range(0, 4):
diff --git a/tests/golden/cases/loop_helper_caller.py b/tests/golden/cases/loop_helper_caller.py
index 666b40f..67a2f2f 100644
--- a/tests/golden/cases/loop_helper_caller.py
+++ b/tests/golden/cases/loop_helper_caller.py
@@ -8,7 +8,7 @@ def loop_helper_caller(x):
             s = s + helper_inner(x)
         return s
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     s = 0.0
     s_1 = s
     for i in range(0, 3):
diff --git a/tests/golden/cases/loop_start_stop_step.py b/tests/golden/cases/loop_start_stop_step.py
index f9a9461..f1e4f2a 100644
--- a/tests/golden/cases/loop_start_stop_step.py
+++ b/tests/golden/cases/loop_start_stop_step.py
@@ -8,7 +8,7 @@ def loop_start_stop_step(x):
             s = s * i
         return s + x
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     s = 1.0
     s_1 = s
     for i in range(2, 8, 3):
diff --git a/tests/golden/cases/loop_two_carried.py b/tests/golden/cases/loop_two_carried.py
index 0bf172d..50617f1 100644
--- a/tests/golden/cases/loop_two_carried.py
+++ b/tests/golden/cases/loop_two_carried.py
@@ -10,7 +10,7 @@ def loop_two_carried(x):
             b = b * 2.0
         return a + b
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     a = 0.0
     b = 1.0
     a_1 = a
diff --git a/tests/golden/cases/loop_with_branch.py b/tests/golden/cases/loop_with_branch.py
index 987b845..2594481 100644
--- a/tests/golden/cases/loop_with_branch.py
+++ b/tests/golden/cases/loop_with_branch.py
@@ -11,7 +11,7 @@ def loop_with_branch(x):
                 s = s - 1.0
         return s
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     s = 0.0
     s_1 = s
     for i in range(0, 4):
diff --git a/tests/golden/cases/math_calls.py b/tests/golden/cases/math_calls.py
index f0f04f2..52038db 100644
--- a/tests/golden/cases/math_calls.py
+++ b/tests/golden/cases/math_calls.py
@@ -5,7 +5,7 @@ def math_calls_vec(x):
     def math_calls(x):
         return math.sqrt(x) + math.exp(x) + math.log1p(x) + math.sin(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.float64)
     return xp.sqrt(t_2) + xp.exp(t_2) + xp.log1p(t_2) + xp.sin(t_2)
diff --git a/tests/golden/cases/math_floor_ceil.py b/tests/golden/cases/math_floor_ceil.py
index 495a20a..b853a37 100644
--- a/tests/golden/cases/math_floor_ceil.py
+++ b/tests/golden/cases/math_floor_ceil.py
@@ -5,7 +5,7 @@ def math_floor_ceil_vec(x):
     def math_floor_ceil(x):
         return math.floor(x) + math.ceil(x)
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.float64)
     return xp.floor(t_2) + xp.ceil(t_2)
diff --git a/tests/golden/cases/math_two_arg.py b/tests/golden/cases/math_two_arg.py
index 07c7397..63441ed 100644
--- a/tests/golden/cases/math_two_arg.py
+++ b/tests/golden/cases/math_two_arg.py
@@ -5,7 +5,7 @@ def math_two_arg_vec(x, y):
     def math_two_arg(x, y):
         return math.atan2(x, y) + math.hypot(x, y) + math.copysign(x, y)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     t_1 = xp.asarray(x)
     t_2 = xp.asarray(y)
     t_3 = xp.astype(t_1, xp.float64)
diff --git a/tests/golden/cases/merge_with_prior.py b/tests/golden/cases/merge_with_prior.py
index 90992f0..3909dfc 100644
--- a/tests/golden/cases/merge_with_prior.py
+++ b/tests/golden/cases/merge_with_prior.py
@@ -8,7 +8,7 @@ def merge_with_prior(x):
             r = x
         return r
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     r = 0.0
     r_1 = x
     r_2 = xp.where(x > 0, r_1, r)
diff --git a/tests/golden/cases/minmax.py b/tests/golden/cases/minmax.py
index 27af9fb..9e0642c 100644
--- a/tests/golden/cases/minmax.py
+++ b/tests/golden/cases/minmax.py
@@ -5,7 +5,7 @@ def minmax_vec(x, y):
     def minmax(x, y):
         return min(x, y) + max(x, y, 3.0)
     """
-    xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x, y)
     t_1 = xp.asarray(x)
     t_2 = xp.asarray(y)
     return _vec_arith(xp, 1, _vec_minmax(xp, True, t_1, t_2), _vec_minmax(xp, False, t_1, t_2, 3.0))
diff --git a/tests/golden/cases/nested_early_returns.py b/tests/golden/cases/nested_early_returns.py
index c8e7a61..903db7a 100644
--- a/tests/golden/cases/nested_early_returns.py
+++ b/tests/golden/cases/nested_early_returns.py
@@ -9,5 +9,5 @@ def nested_early_returns(x):
             return 2.0
         return 3.0
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.where(x > 0, xp.where(x > 10, 1.0, 2.0), 3.0)
diff --git a/tests/golden/cases/nested_loops.py b/tests/golden/cases/nested_loops.py
index 3092987..4bcfebe 100644
--- a/tests/golden/cases/nested_loops.py
+++ b/tests/golden/cases/nested_loops.py
@@ -9,7 +9,7 @@ def nested_loops(x):
                 s = s + x
         return s
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     s = 0.0
     s_1 = s
     for i in range(0, 3):
diff --git a/tests/golden/cases/piecewise.py b/tests/golden/cases/piecewise.py
index 153885f..33b08da 100644
--- a/tests/golden/cases/piecewise.py
+++ b/tests/golden/cases/piecewise.py
@@ -9,5 +9,5 @@ def piecewise(x):
             return 1.0
         return x
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.where(x < -1, -1.0, xp.where(x > 1, 1.0, x))
diff --git a/tests/golden/cases/pow_op.py b/tests/golden/cases/pow_op.py
index 9872b9f..222a4a1 100644
--- a/tests/golden/cases/pow_op.py
+++ b/tests/golden/cases/pow_op.py
@@ -5,5 +5,5 @@ def pow_op_vec(x):
     def pow_op(x):
         return x**2 + 2.0**x
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.pow(x, 2) + xp.pow(2.0, x)
diff --git a/tests/golden/cases/psi.py b/tests/golden/cases/psi.py
index d4aa00c..afa870f 100644
--- a/tests/golden/cases/psi.py
+++ b/tests/golden/cases/psi.py
@@ -7,5 +7,5 @@ def psi(x):
             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(xp.astype(xp.asarray(-x), xp.float64)))
diff --git a/tests/golden/cases/relu.py b/tests/golden/cases/relu.py
index 54e4361..8cb6178 100644
--- a/tests/golden/cases/relu.py
+++ b/tests/golden/cases/relu.py
@@ -7,5 +7,5 @@ def relu(x):
             return x
         return 0.0
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.where(x > 0, x, 0.0)
diff --git a/tests/golden/cases/relu_with_else.py b/tests/golden/cases/relu_with_else.py
index 7eb914f..b2fac47 100644
--- a/tests/golden/cases/relu_with_else.py
+++ b/tests/golden/cases/relu_with_else.py
@@ -9,7 +9,7 @@ def relu_with_else(x):
             r = 0.0
         return r
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     r_1 = x
     r_2 = 0.0
     r = xp.where(x > 0, r_1, r_2)
diff --git a/tests/golden/cases/reserved_param.py b/tests/golden/cases/reserved_param.py
index e1d1464..987a18f 100644
--- a/tests/golden/cases/reserved_param.py
+++ b/tests/golden/cases/reserved_param.py
@@ -5,5 +5,5 @@ def reserved_param_vec(xp):
     def reserved_param(xp):
         return xp + 1
     """
-    xp_1 = array_namespace(*[a for a in (xp,) if hasattr(a, '__array_namespace__')])
+    xp_1 = array_namespace(xp)
     return xp + 1
diff --git a/tests/golden/cases/ssa_rebind.py b/tests/golden/cases/ssa_rebind.py
index e9d82fa..34fb44b 100644
--- a/tests/golden/cases/ssa_rebind.py
+++ b/tests/golden/cases/ssa_rebind.py
@@ -8,7 +8,7 @@ def ssa_rebind(x):
         y += 3
         return y
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     y = x + 1
     y_1 = y * 2
     y_2 = y_1 + 3
diff --git a/tests/golden/cases/ternary.py b/tests/golden/cases/ternary.py
index 8e0f2a6..4a8e2ee 100644
--- a/tests/golden/cases/ternary.py
+++ b/tests/golden/cases/ternary.py
@@ -5,5 +5,5 @@ def ternary_vec(x):
     def ternary(x):
         return x * 2 if x > 0 else x / 2
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return xp.where(x > 0, x * 2, x / 2)
diff --git a/tests/golden/cases/unary_ops.py b/tests/golden/cases/unary_ops.py
index a69b032..746977c 100644
--- a/tests/golden/cases/unary_ops.py
+++ b/tests/golden/cases/unary_ops.py
@@ -5,5 +5,5 @@ def unary_ops_vec(x):
     def unary_ops(x):
         return -x + +x - ~x
     """
-    xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
+    xp = array_namespace(x)
     return -x + +x - ~x
diff --git a/tests/unit/test_codegen.py b/tests/unit/test_codegen.py
index 16dd6f1..3d52aa2 100644
--- a/tests/unit/test_codegen.py
+++ b/tests/unit/test_codegen.py
@@ -53,7 +53,7 @@ def test_module_shape() -> None:
     assert "def f_vec(x):" in src
     assert '"""\n    def f(x):\n        return x\n    """' in src
     assert "xp = array_namespace(" in src
-    assert "hasattr(a, '__array_namespace__')" in src
+    assert "hasattr" not in src
 
 
 def test_operators_stay_operators() -> None:
@@ -171,7 +171,7 @@ def test_hidden_params_kwonly_none() -> None:
     )
     src = generate_source(lowered, lowered.program)
     assert "def f_vec(x, *, ARR=None):" in src
-    assert "hasattr(a, '__array_namespace__')" in src
+    assert "xp = array_namespace(x, ARR)" in src
 
 
 def test_posonly_marker() -> None:

From c163eb8e92d276a493c5d0f74c0d99b74864530e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= <enriqueg@cscs.ch>
Date: Mon, 5 Oct 2026 21:50:04 +0200
Subject: [PATCH 2/4] compat: delegate array detection to array_api_compat
 (verify/fallback/closures work on torch)

---
 src/array_vectorize/compat.py           | 22 +++++++++++++++++++++-
 src/array_vectorize/frontend/extract.py |  3 ++-
 src/array_vectorize/verify.py           | 12 +++++++++---
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/array_vectorize/compat.py b/src/array_vectorize/compat.py
index a36874e..dea8b73 100644
--- a/src/array_vectorize/compat.py
+++ b/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:
diff --git a/src/array_vectorize/frontend/extract.py b/src/array_vectorize/frontend/extract.py
index 80b1c1b..dddf5e4 100644
--- a/src/array_vectorize/frontend/extract.py
+++ b/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
diff --git a/src/array_vectorize/verify.py b/src/array_vectorize/verify.py
index 44ef8cc..6c49044 100644
--- a/src/array_vectorize/verify.py
+++ b/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

From 23211631f4b2ab5e1378e050ffe97a19ec5337ef Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= <enriqueg@cscs.ch>
Date: Mon, 5 Oct 2026 21:50:13 +0200
Subject: [PATCH 3/4] tests: jax.jit / torch.compile compatibility + torch
 detection regression tests

---
 pyproject.toml                               |   7 +
 tests/backends/test_jax_torch_compilation.py | 255 +++++++++++++++++++
 2 files changed, 262 insertions(+)
 create mode 100644 tests/backends/test_jax_torch_compilation.py

diff --git a/pyproject.toml b/pyproject.toml
index 42933f6..5df7438 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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",
diff --git a/tests/backends/test_jax_torch_compilation.py b/tests/backends/test_jax_torch_compilation.py
new file mode 100644
index 0000000..101358c
--- /dev/null
+++ b/tests/backends/test_jax_torch_compilation.py
@@ -0,0 +1,255 @@
+"""jax.jit / torch.compile compatibility + backend-detection regression tests.
+
+JAX and PyTorch trace and introspect the Python they compile — ``jax.jit``
+runs the generated body under tracers (through the namespace-detection
+preamble), ``torch.compile`` traces it with Dynamo — so the generated code
+must be valid, inspectable, traceable Python. Every corpus function below
+is exercised in four modes against its eager result:
+
+- ``jax.jit`` with the namespace detected from the arguments
+- ``jax.jit`` with the namespace pinned (``vectorize(fn, namespace=jax.numpy)``;
+  the pinned generated source has no detection code at all)
+- ``torch.compile`` with the namespace detected from the arguments
+- ``torch.compile`` with the namespace pinned to ``array_api_compat.torch``
+  (what ``array_namespace()`` returns for torch tensors)
+
+Beyond compilation, this file regressions the backend detection delegated
+to array-api-compat (``array_namespace``): ``verify=``, ``fallback=True``
+and closure captures must recognize torch tensors, whose class has never
+exposed ``__array_namespace__``.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+from support import make_fn, make_module
+
+from array_vectorize import vectorize
+
+spec = importlib.util.spec_from_file_location(
+    "vec_corpus_c", Path(__file__).parent.parent / "corpus.py"
+)
+assert spec is not None and spec.loader is not None
+CORPUS = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(CORPUS)
+
+X = [-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]
+Y = [3.0, 1.5, 1.0, 0.5, -1.0, -3.0]
+# closure_array's captured ARR is [1.0, 2.0, 3.0]: a length-3 input
+# broadcasts against it, and ARR is passed explicitly as a backend array —
+# detected mode needs a single namespace among the arguments, and pinned
+# torch must not hit NumPy's deprecated __array_wrap__ path.
+X3 = [-2.0, 0.5, 3.0]
+ARR = [1.0, 2.0, 3.0]
+
+TWO_ARG = frozenset({"add", "math_two_arg", "minmax"})
+
+NAMES = [
+    "add",
+    "relu",
+    "relu_with_else",
+    "psi",
+    "piecewise",
+    "nested_early_returns",
+    "math_calls",
+    "math_two_arg",
+    "from_import",
+    "closure_array",
+    "closure_scalar",
+    "loop_accumulate",
+    "loop_start_stop_step",
+    "loop_two_carried",
+    "loop_with_branch",
+    "nested_loops",
+    "kwonly_defaults",
+    "helper_outer",
+    "loop_helper_caller",
+    "ternary",
+    "minmax",
+    "cse_opportunity",
+]
+
+
+def _call_args(
+    name: str, conv: Callable[[Sequence[float]], Any]
+) -> tuple[tuple[Any, ...], dict[str, Any]]:
+    """Positional args (plus the closure kwarg) for ``CORPUS.<name>``."""
+    if name == "closure_array":
+        return (conv(X3),), {"ARR": conv(ARR)}
+    if name in TWO_ARG:
+        return (conv(X), conv(Y)), {}
+    return (conv(X),), {}
+
+
+def _torch_conv(vals: Sequence[float]) -> Any:
+    import torch
+
+    return torch.tensor(vals, dtype=torch.float64)
+
+
+# ---- jax.jit
+
+
+@pytest.mark.parametrize("name", NAMES)
+def test_jax_jit_detected_matches_eager(name: str) -> None:
+    jax = pytest.importorskip("jax")
+    # the generated math-call lowering casts through float64 (Python float
+    # semantics); without x64 jax warns on every astype and truncates
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    fn = getattr(CORPUS, name)
+    args, kwargs = _call_args(name, jnp.asarray)
+    eager = vectorize(fn)(*args, **kwargs)
+    got = jax.jit(vectorize(fn))(*args, **kwargs)
+    assert np.allclose(np.asarray(got), np.asarray(eager), equal_nan=True)
+
+
+@pytest.mark.parametrize("name", NAMES)
+def test_jax_jit_pinned_matches_eager(name: str) -> None:
+    jax = pytest.importorskip("jax")
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    fn = getattr(CORPUS, name)
+    args, kwargs = _call_args(name, jnp.asarray)
+    eager = vectorize(fn)(*args, **kwargs)
+    got = jax.jit(vectorize(fn, namespace=jnp))(*args, **kwargs)
+    assert np.allclose(np.asarray(got), np.asarray(eager), equal_nan=True)
+
+
+# ---- torch.compile
+#
+# NOTE: detected-mode torch used to fail with "TypeError: array_namespace
+# requires at least one non-scalar array input" — the generated preamble
+# filtered arguments with ``hasattr(a, '__array_namespace__')``, which
+# torch.Tensor has never exposed (any 2.x). The preamble now passes the
+# arguments straight to array_namespace() (array-api-compat >= 1.10 skips
+# Python scalars/None and dispatches torch tensors), so all four modes
+# below must pass.
+
+
+@pytest.mark.timeout(600)
+@pytest.mark.parametrize("name", NAMES)
+def test_torch_compile_detected_matches_eager(name: str) -> None:
+    pytest.importorskip("torch")
+    import torch
+
+    fn = getattr(CORPUS, name)
+    args, kwargs = _call_args(name, _torch_conv)
+    eager = vectorize(fn)(*args, **kwargs)
+    got = torch.compile(vectorize(fn))(*args, **kwargs)
+    assert torch.allclose(got, eager, equal_nan=True)
+
+
+@pytest.mark.timeout(600)
+@pytest.mark.parametrize("name", NAMES)
+def test_torch_compile_pinned_matches_eager(name: str) -> None:
+    pytest.importorskip("torch")
+    import array_api_compat.torch
+    import torch
+
+    fn = getattr(CORPUS, name)
+    args, kwargs = _call_args(name, _torch_conv)
+    eager = vectorize(fn, namespace=array_api_compat.torch)(*args, **kwargs)
+    got = torch.compile(vectorize(fn, namespace=array_api_compat.torch))(*args, **kwargs)
+    assert torch.allclose(got, eager, equal_nan=True)
+
+
+# ---- backend-detection regression tests (delegated to array_namespace)
+#
+# verify=, fallback=True and closure captures used to detect arrays with
+# a homegrown ``hasattr(value, '__array_namespace__')`` check, which
+# torch.Tensor has never satisfied — these paths now delegate to
+# array_api_compat's array_namespace and must work on torch tensors.
+
+
+def test_verify_with_torch_args() -> None:
+    pytest.importorskip("torch")
+    import torch
+
+    t = torch.tensor(X, dtype=torch.float64)
+    # previously: VectorizationError("verify= needs at least one Array API
+    # array in the example inputs")
+    vec = vectorize(CORPUS.psi, verify=(t,))
+    got = vec(t)
+    expected = torch.where(t < 0, torch.tensor(0.0), t * torch.exp(-t))
+    assert torch.allclose(got, expected)
+
+
+def test_fallback_with_torch_tensors() -> None:
+    pytest.importorskip("torch")
+    import torch
+
+    fn = make_fn("    while x > 0:\n        x = x - 1\n    return x", defaults="x")
+    with pytest.warns(UserWarning, match="falling back"):
+        vec = vectorize(fn, fallback=True)
+    # previously: TypeError("vectorized fallback ... requires at least one
+    # Array API array argument")
+    t = torch.tensor([5.0, -2.0], dtype=torch.float64)
+    got = vec(t)
+    assert isinstance(got, torch.Tensor)
+    # the element loop re-wraps Python scalars via xp.asarray, which on
+    # torch lands on the default float32 — values, not dtype, are the
+    # fallback contract (mirrors the numpy fallback tests' list compare)
+    assert torch.allclose(got.double(), torch.tensor([0.0, -2.0], dtype=torch.float64))
+
+
+def test_fallback_pinned_with_torch_tensors() -> None:
+    pytest.importorskip("torch")
+    import array_api_compat.torch
+    import torch
+
+    fn = make_fn("    while x > 0:\n        x = x - 1\n    return x", defaults="x")
+    with pytest.warns(UserWarning, match="falling back"):
+        vec = vectorize(fn, fallback=True, namespace=array_api_compat.torch)
+    t = torch.tensor([5.0, -2.0], dtype=torch.float64)
+    got = vec(t)
+    assert isinstance(got, torch.Tensor)
+    assert torch.allclose(got.double(), torch.tensor([0.0, -2.0], dtype=torch.float64))
+
+
+def test_closure_captured_torch_tensor() -> None:
+    pytest.importorskip("torch")
+    import torch
+
+    mod = make_module(
+        "import torch\n"
+        "T = torch.asarray([1.0, 2.0, 3.0], dtype=torch.float64)\n"
+        "\n"
+        "\n"
+        "def subject(x):\n"
+        "    return x + T\n"
+    )
+    # previously: VectorizationError (closure/global 'T' has unsupported
+    # type 'Tensor'); now a hidden kw-only closure param like numpy captures
+    vec = vectorize(mod.subject)
+    assert "def subject_vec(x, *, T=None):" in vec.source
+    assert vec.__kwdefaults__["T"] is mod.T
+    t = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float64)
+    got = vec(t)
+    assert torch.allclose(got, torch.tensor([11.0, 22.0, 33.0], dtype=torch.float64))
+
+
+def test_closure_captured_jax_array() -> None:
+    jax = pytest.importorskip("jax")
+    jax.config.update("jax_enable_x64", True)
+    mod = make_module(
+        "import jax.numpy as jnp\n"
+        "T = jnp.asarray([1.0, 2.0, 3.0])\n"
+        "\n"
+        "\n"
+        "def subject(x):\n"
+        "    return x + T\n"
+    )
+    vec = vectorize(mod.subject)
+    assert "def subject_vec(x, *, T=None):" in vec.source
+    assert vec.__kwdefaults__["T"] is mod.T
+    got = vec(mod.T * 10.0)
+    assert np.allclose(np.asarray(got), [11.0, 22.0, 33.0])

From a4f3d6b1fa292e7f3338469e53baa6fdd931b8e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= <enriqueg@cscs.ch>
Date: Mon, 5 Oct 2026 21:50:13 +0200
Subject: [PATCH 4/4] ci: compile job + make backends target

---
 .github/workflows/ci.yml | 17 +++++++++++++++++
 Makefile                 |  5 ++++-
 docs/development.md      |  6 ++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 05fd273..bda36d0 100644
--- a/.github/workflows/ci.yml
+++ b/.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:
diff --git a/Makefile b/Makefile
index 5a02e20..b3ee5b3 100644
--- a/Makefile
+++ b/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
 
diff --git a/docs/development.md b/docs/development.md
index 8fc6a8f..36ae56b 100644
--- a/docs/development.md
+++ b/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