ENH: fix is_jax_array under torch.compile - #464
Conversation
Sorry, something went wrong.
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)) ```
|
Is there an MWE, ideally as a test or as an example for the added comment? attrs on empty tuples is a tad cryptic, and I'm not sure how to trigger it---or check if the whole _is_zero_grad_array can be removed, as the comment above indicates will be possible "at some point". |
Sorry, something went wrong.
|
calling is_jax_array on a torch array while running under torch.compile(fullgraph=True) should reproduce. |
Sorry, something went wrong.
|
Great. Add an otherwise failing code snippet to the comment? |
Sorry, something went wrong.
|
CI failures in test_trace are unrelated, data-apis/array-api-tests#455
Does not repro locally: In [11]: import torch In [12]: import array_api_compat as aac In [13]: _is_zero_grad = aac.common._helpers._is_jax_zero_gradient_array In [14]: cfunc = torch.compile(lambda x: _is_zero_grad(x), fullgraph=True) In [15]: cfunc(torch.ones(3)) Out[15]: False In [16]: torch.__version__ Out[16]: '2.7.1' In [17]: aac.__version__ Out[17]: '1.16.0.dev0' |
Sorry, something went wrong.
|
Interesting, I'll have to take a closer look. |
Sorry, something went wrong.
This fixes errors like the following:
cc @crusaderky for vis