Skip to content

[codex] Add JAX oracle witnesses and unified notes#19

Merged
shinaoka merged 14 commits intomainfrom
feat/unified-jax-pytorch-oracles
Apr 6, 2026
Merged

[codex] Add JAX oracle witnesses and unified notes#19
shinaoka merged 14 commits intomainfrom
feat/unified-jax-pytorch-oracles

Conversation

@shinaoka
Copy link
Copy Markdown
Member

@shinaoka shinaoka commented Apr 5, 2026

Summary

Add the initial unified JAX/PyTorch oracle expansion on top of the existing tensor-ad-oracles layout.

This PR:

  • adds JAX witness materialization for published abs/identity and exp/identity cases, including jvp, vjp, linearization, transpose, and adjoint_check
  • backfills the published cases/abs/identity.jsonl and cases/exp/identity.jsonl records with jax_ref
  • updates replay and regeneration integrity checks so CI validates the JAX witness surface alongside the existing PyTorch oracles
  • unifies the math-note structure across published notes and adds README clarification that validate_schema.py is a maintainer/CI integrity tool while published JSONL files already store materialized numeric payloads

Why

The branch extends the current oracle database instead of splitting a separate JAX repository, while making the JAX convention explicit in both the machine-readable witnesses and the math-note corpus.

Validation

  • uv run python -m unittest tests.test_repo_config tests.test_scripts tests.test_materialize tests.test_pytorch_v1 tests.test_validator_encoding tests.test_db_replay tests.test_jax_v1 tests.test_runtime_jax tests.test_math_notes_structure tests.test_math_registry -v
  • uv run python scripts/verify_cases.py
  • uv run python scripts/check_replay.py
  • uv run python scripts/check_regeneration.py

Notes

  • uv run python scripts/validate_schema.py was measured separately on the current branch and completed successfully in about 6 minutes for 9580 records.

@shinaoka shinaoka marked this pull request as ready for review April 6, 2026 00:50
@shinaoka shinaoka merged commit aed534a into main Apr 6, 2026
2 checks passed
@shinaoka shinaoka deleted the feat/unified-jax-pytorch-oracles branch April 6, 2026 01:16
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