feat(detection): add handler for remat2 primitive

[adrhill]

Summary

• Add support for jax.checkpoint / jax.remat by handling the remat2 primitive
• The primitive uses the same "jaxpr" parameter key as jit/pjit, so it reuses the existing prop_closed_jaxpr implementation

Test plan

• test_remat2_checkpoint — sparsity detection through checkpoint
• test_remat2_closure_captured_index — closure-captured indices resolve precisely
• test_remat2_decompression — full Jacobian computation works

Closes #118

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.14%. Comparing base (b5b3c54) to head (0a8be0b).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #141   +/-   ##
=======================================
  Coverage   93.14%   93.14%           
=======================================
  Files          53       53           
  Lines        3518     3518           
=======================================
  Hits         3277     3277           
  Misses        241      241           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Adds sparsity-detection support for jax.checkpoint / jax.remat by handling the remat2 JAX primitive as a nested-jaxpr call, consistent with how jit/pjit/named_call are already handled in the interpreter.

Changes:

• Dispatch remat2 equations through the existing prop_closed_jaxpr(..., param_key="jaxpr") path.
• Add interpreter tests covering checkpoint/remat sparsity, closure-captured indices, and full Jacobian decompression through jax.checkpoint.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
src/asdex/detection/_interpret/__init__.py	Treats remat2 like other closed-jaxpr call primitives by reusing the jaxpr-param propagation handler.
tests/_interpret/test_nested_jaxpr.py	Adds regression tests to ensure checkpoint/remat traces correctly for sparsity and decompression, including precise gather tracking with closure-captured indices.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.