Skip to content

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

Merged
egparedes merged 4 commits into
mainfrom
backend-compile-tests
Oct 5, 2026
Merged

egparedes merged 4 commits into
mainfrom
backend-compile-tests

Conversation

@egparedes

Copy link
Copy Markdown
Contributor

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__'), but torch.Tensor has never exposed that attribute in any 2.x release — torch tensors were dropped like scalars and array_namespace() raised TypeError. The same homegrown duck-check also broke verify=, 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

  • Generated preamble: xp = array_namespace(x, y) — arguments passed straight through (compat ignores Python scalars/None since 1.10 and dispatches torch/CuPy). Floor bumped to array-api-compat>=1.10; 41 goldens regenerated (namespace line only).
  • _is_array (verify + fallback): fully delegated to array_namespace via try/except — no homegrown partial implementations of the standard's detection machinery.
  • Closure classification (extract.py): captured torch tensors now become hidden closure params (previously misclassified). Verified no numpy reclassifications occur.
  • Bonus fix: verify= built its expected side with xp.asarray(list), which torch defaults to float32 — now explicit dtype=xp.float64 (backend-neutral).

Tests (93, zero skips in the compile job)

  • 88 compilation tests: 22 corpus functions (branches, loops, closures, helpers, kw-only defaults) × jax.jit detected/pinned × torch.compile detected/pinned, compared against eager results.
  • 5 detection regression tests: verify-with-torch, fallback-with-torch, pinned fallback, closure-captured torch tensor, jax closure twin.

CI / infra

  • New compile job (CPU-only torch wheel, jax; import guard so skips can't hide failures), make backends target, backends optional 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 check 804 passed / 96% coverage; make docs clean; compat==1.10 floor verified green (711 passed, one over-specified error-message regex widened). Zero graph breaks in torch.compile pinned mode.

@egparedes
egparedes merged commit f0354b4 into main Oct 5, 2026
17 of 24 checks passed
@egparedes
egparedes deleted the backend-compile-tests branch October 6, 2026 05:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant