From 1f66a47664f0cfc063c93cde6437be52706864fb Mon Sep 17 00:00:00 2001
From: Lucas Colley <lucas.colley8@gmail.com>
Date: Sun, 16 Aug 2026 15:03:11 +0100
Subject: [PATCH] BUG: fix `is_jax_array` under `torch.compile`

This fixes errors like the following:
```
E               torch._dynamo.exc.UserError:
GetAttrVariable(TupleVariable(length=0), dtype) has no type
E               For more information about this error, see:
https://pytorch.org/docs/main/generated/exportdb/index.html#unknown-python-type
E
E               from user code:
E                  File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_lib/_helpers.py",
line 612, in inner
E                   res = func(*args, **kwargs)  # pyright:
ignore[reportCallIssue]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/tests/main/test_at.py",
line 47, in at_op
E                   return meth(y, copy=copy, xp=xp)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 374, in add
E                   return self._op(_AtOp.ADD, operator.iadd,
operator.add, y, copy=copy, xp=xp)  # pyright:
ignore[reportUnknownArgumentType]
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/src/array_api_extra/_at.py",
line 289, in _op
E                   (_compat.is_dask_array(idx) or
_compat.is_jax_array(idx))
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 253, in is_jax_array
E                   or _is_jax_zero_gradient_array(x)
E                 File
"/Users/lucascolley/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-backends/lib/python3.14/site-packages/array_api_compat/common/_helpers.py",
line 84, in _is_jax_zero_gradient_array
E                   cls = cast(Hashable, type(dtype))
```
---
 src/array_api_compat/common/_helpers.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/array_api_compat/common/_helpers.py b/src/array_api_compat/common/_helpers.py
index d5342658..5c0574ef 100644
--- a/src/array_api_compat/common/_helpers.py
+++ b/src/array_api_compat/common/_helpers.py
@@ -80,6 +80,13 @@ def _is_jax_zero_gradient_array(x: object) -> TypeGuard[_ZeroGradientArray]:
         dtype = x.dtype  # type: ignore[attr-defined]
     except AttributeError:
         return False
+
+    torch = sys.modules.get("torch")
+    if torch is not None and torch.compiler.is_compiling():
+        # Under torch.compile tracing, `type(dtype)` can't be modeled by Dynamo
+        # for objects it hasn't fully specialized (e.g. attrs on empty tuples).
+        return False
+
     cls = cast(Hashable, type(dtype))
     if not _issubclass_fast(cls, "numpy.dtypes", "VoidDType"):
         return False