Repository navigation
jax.jit / torch.compile tests + fix torch detected mode (delegate detection to array-api-compat) - #6
Merged
Merged
Conversation
…/closures work on torch)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds framework-compilation tests for generated code — and fixes the real bug they uncovered: torch tensors never worked in detected mode.
The bug
The generated preamble filtered arguments with
hasattr(a, '__array_namespace__'), buttorch.Tensorhas never exposed that attribute in any 2.x release — torch tensors were dropped like scalars andarray_namespace()raisedTypeError. The same homegrown duck-check also brokeverify=, the fallback wrapper, and closure captures for torch. This contradicted the README/docs claims of torch support; it was invisible because torch was never installed in CI.The fix — delegate ALL array detection to array-api-compat
xp = array_namespace(x, y)— arguments passed straight through (compat ignores Python scalars/None since 1.10 and dispatches torch/CuPy). Floor bumped toarray-api-compat>=1.10; 41 goldens regenerated (namespace line only)._is_array(verify + fallback): fully delegated toarray_namespacevia try/except — no homegrown partial implementations of the standard's detection machinery.extract.py): captured torch tensors now become hidden closure params (previously misclassified). Verified no numpy reclassifications occur.verify=built its expected side withxp.asarray(list), which torch defaults to float32 — now explicitdtype=xp.float64(backend-neutral).Tests (93, zero skips in the compile job)
jax.jitdetected/pinned ×torch.compiledetected/pinned, compared against eager results.CI / infra
compilejob (CPU-only torch wheel, jax; import guard so skips can't hide failures),make backendstarget,backendsoptional dependency group, docs note (torch.compile of detected-mode functions needs compat >= 1.15).Verification
93/93 framework tests (jax 0.11.2, torch 2.14.1+cpu);
make check804 passed / 96% coverage;make docsclean; compat==1.10 floor verified green (711 passed, one over-specified error-message regex widened). Zero graph breaks in torch.compile pinned mode.