Skip to content

M13: EXLA gradient conformance - #23

Merged
ausimian merged 1 commit into
mainfrom
m13-exla-grad-conformance
Apr 16, 2026
Merged

M13: EXLA gradient conformance#23
ausimian merged 1 commit into
mainfrom
m13-exla-grad-conformance

Conversation

@ausimian

Copy link
Copy Markdown
Owner

Summary

  • Add EXLA (XLA CPU backend) as a third gradient oracle, independent of both Nx.BinaryBackend and MLX
  • Extract shared defn grad functions into Emily.GradZoo (eliminates duplication across test files)
  • 9 tests: 8 zoo functions + 1 transformer-block training step, compared against checked-in EXLA golden values with per-function tolerance tables
  • Golden generator is a single Mix.install script: elixir bench/exla_golden_gen.exs
  • CUDA conformance deferred to post-1.0; XLA-CPU provides the independent lowering path

Test plan

  • mix test test/emily/grad/ — all 22 tests pass (equivalence, finite-diff, EXLA oracle)
  • mix test --only grad_conformance — 9 EXLA oracle tests pass in isolation
  • mix precommit — full suite green

Note: pre-existing flaky test in transform_test.exs:105 (seed-dependent, 15% tolerance boundary) — documented in PLAN.md, not introduced by this PR.

Add a third gradient oracle — EXLA (XLA CPU backend) — to catch bugs
where Emily and BinaryBackend agree on the wrong gradient (they share
the same Nx.Defn.grad lowering). Eight zoo functions plus a full
transformer-block training step are tested against checked-in EXLA
golden values with per-function tolerance tables.

- Extract shared defn functions into Emily.GradZoo
- Generate golden values via bench/exla_golden_gen.exs (EXLA 0.11.0 CPU)
- Test harness at test/emily/grad/exla_oracle_test.exs
- CUDA conformance deferred to post-1.0
@ausimian
ausimian merged commit d7e3ea2 into main Apr 16, 2026
1 check passed
@ausimian
ausimian deleted the m13-exla-grad-conformance branch April 16, 2026 07:11
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