JAX 0.11 compatibility (flattree scan port) + rename CLAUDE.md→AGENTS.md#140
Merged
Conversation
Track the agent instruction file as AGENTS.md (previously the gitignored, untracked CLAUDE.md). Remove the /AGENTS.md ignore entry so it is committed; keep /CLAUDE.md ignored so any local CLAUDE.md stays private.
JAX 0.11 merged ClosedJaxpr into a unified Jaxpr whose first parameter is constvars, dropping the jaxpr= keyword. Pass jaxpr/consts positionally in ModuleInfo.add_jaxpr_outs so it works on old and new JAX, matching the other ClosedJaxpr call sites (fixes TypeError in the compile tests). jax.core also dropped new_jaxpr_eqn in 0.11; annotate the older-JAX import fallback with a type: ignore so mypy passes against new JAX. Add jax 0.10.0 to the CI test matrix.
JAX 0.11 refactored the scan primitive: num_consts/num_carry/linear/ _split_transpose params were removed and the (consts, carry, xs) input split and (carry, ys) output split are now encoded in ft_in/ft_out flattrees. braintrace read num_consts/num_carry directly and rebuilt one scan eqn (add_scan_ys), which broke on 0.11 with KeyError: 'num_consts'. Add two capability-detecting shims in _compatible_imports (detect by which params exist, so the same code path serves every jax >=0.8.0): - scan_num_consts_carry(eqn): read params on jax <0.11, else derive from ft_in.unpack(). - scan_params_add_ys(params, n_extra): identity on jax <0.11 (num_ys is implicit), else grow ft_out's ys component by n_extra leaves. Route all num_consts/num_carry reads (canonicalize, hidden_group, scan_descent) through the shim, and rebuild ft_out in add_scan_ys. while/ cond params are unchanged in 0.11 and untouched. Full suite green on jax 0.11 (2128 passed) and mypy clean; CI matrix covers 0.8/0.9/0.10 (old encoding) and latest (flattree).
Simplify the spec rule to 'Write specs under docs/specs before implementation' and fold the brainstate.random rule into the working agreement list.
Reviewer's GuideMakes braintrace compatible with JAX 0.11’s scan flattree refactor and ClosedJaxpr merge while preserving support for older JAX versions via capability-detecting shims, extends tests and CI to cover the new behavior, and renames the contributor guidance document to AGENTS.md as a tracked file. File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- In
_compatible_imports.scan_params_add_ys, consider caching thejax._src.flattreeimport at module scope instead of inside the function to avoid repeated imports on hot paths. - The new helpers (
scan_num_consts_carry,scan_params_add_ys) currently use very loose types (Dict,JaxprEqn-like); tightening their type hints and explicitly documenting expected param keys would make misuse easier to detect and maintainers clearer on supported shapes.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `_compatible_imports.scan_params_add_ys`, consider caching the `jax._src.flattree` import at module scope instead of inside the function to avoid repeated imports on hot paths.
- The new helpers (`scan_num_consts_carry`, `scan_params_add_ys`) currently use very loose types (`Dict`, `JaxprEqn`-like); tightening their type hints and explicitly documenting expected param keys would make misuse easier to detect and maintainers clearer on supported shapes.
## Individual Comments
### Comment 1
<location path="AGENTS.md" line_range="20" />
<code_context>
+3. After writing code, list edge cases + suggest test cases.
+4. Bug? Write a test that reproduces it, then fix until the test passes.
+5. Every correction: reflect on the mistake, plan to avoid repeating it.
+6. All updates must be happened on the worktree branch, not main.
+7. Write specs under `docs/specs` before implementation.
+8. Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines.
</code_context>
<issue_to_address>
**issue (typo):** The phrase "must be happened" is ungrammatical; consider "must happen" or "must be made".
Consider updating this sentence to: "All updates must happen on the worktree branch, not main." or "All updates must be made on the worktree branch, not main."
```suggestion
6. All updates must happen on the worktree branch, not main.
```
</issue_to_address>
### Comment 2
<location path="AGENTS.md" line_range="22" />
<code_context>
+5. Every correction: reflect on the mistake, plan to avoid repeating it.
+6. All updates must be happened on the worktree branch, not main.
+7. Write specs under `docs/specs` before implementation.
+8. Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines.
+9. Co-locate tests with the code under test: each module `foo.py` has its tests in a sibling `foo_test.py` (suffix style — never a separate `tests/` directory, never the `test_*.py` prefix).
+10. **Never drive a model with a bare Python `for`/`while` loop when it runs repeatedly.** Python loops execute op-by-op (dispatch overhead, no fusion) and trace fresh each step; the `brainstate.transform` primitives lower the whole loop into one compiled XLA program, tracing the body only once. Pick by shape of the work:
</code_context>
<issue_to_address>
**issue (typo):** Missing verb in "Tests should >90% coverage"; consider adding "have" or similar.
Consider "Tests should have >90% coverage" or "Tests should achieve >90% coverage" to fix the grammatical omission.
```suggestion
8. Tests should have >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines.
```
</issue_to_address>
### Comment 3
<location path="AGENTS.md" line_range="146" />
<code_context>
+- JAX as the core computation framework.
+- Additional brain-ecosystem utilities.
+
+Pin exact names/versions in `pyproject.toml` / requirements files,.
+
+## Known limitations
</code_context>
<issue_to_address>
**issue (typo):** There is an extra comma before the period at the end of this sentence.
Please remove the trailing comma so the sentence ends with just "files."
```suggestion
Pin exact names/versions in `pyproject.toml` / requirements files.
```
</issue_to_address>
### Comment 4
<location path="AGENTS.md" line_range="175" />
<code_context>
+
+#### Rules for the Examples section
+
+- Wrap example code in `.. code-block:: python` directive so Sphinx render with syntax highlighting.
+- Prefix every input line with `>>>` (continuation lines with `...`) for `doctest` compatibility.
+- Show expected output on line immediately after statement, **without** prompt prefix.
</code_context>
<issue_to_address>
**issue (typo):** Subject-verb agreement issue: "Sphinx render" should be "Sphinx renders".
Update the bullet to: "Wrap example code in `.. code-block:: python` directive so Sphinx renders with syntax highlighting."
```suggestion
- Wrap example code in `.. code-block:: python` directive so Sphinx renders with syntax highlighting.
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| 3. After writing code, list edge cases + suggest test cases. | ||
| 4. Bug? Write a test that reproduces it, then fix until the test passes. | ||
| 5. Every correction: reflect on the mistake, plan to avoid repeating it. | ||
| 6. All updates must be happened on the worktree branch, not main. |
There was a problem hiding this comment.
issue (typo): The phrase "must be happened" is ungrammatical; consider "must happen" or "must be made".
Consider updating this sentence to: "All updates must happen on the worktree branch, not main." or "All updates must be made on the worktree branch, not main."
Suggested change
| 6. All updates must be happened on the worktree branch, not main. | |
| 6. All updates must happen on the worktree branch, not main. |
| 5. Every correction: reflect on the mistake, plan to avoid repeating it. | ||
| 6. All updates must be happened on the worktree branch, not main. | ||
| 7. Write specs under `docs/specs` before implementation. | ||
| 8. Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines. |
There was a problem hiding this comment.
issue (typo): Missing verb in "Tests should >90% coverage"; consider adding "have" or similar.
Consider "Tests should have >90% coverage" or "Tests should achieve >90% coverage" to fix the grammatical omission.
Suggested change
| 8. Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines. | |
| 8. Tests should have >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines. |
| - JAX as the core computation framework. | ||
| - Additional brain-ecosystem utilities. | ||
|
|
||
| Pin exact names/versions in `pyproject.toml` / requirements files,. |
There was a problem hiding this comment.
issue (typo): There is an extra comma before the period at the end of this sentence.
Please remove the trailing comma so the sentence ends with just "files."
Suggested change
| Pin exact names/versions in `pyproject.toml` / requirements files,. | |
| Pin exact names/versions in `pyproject.toml` / requirements files. |
|
|
||
| #### Rules for the Examples section | ||
|
|
||
| - Wrap example code in `.. code-block:: python` directive so Sphinx render with syntax highlighting. |
There was a problem hiding this comment.
issue (typo): Subject-verb agreement issue: "Sphinx render" should be "Sphinx renders".
Update the bullet to: "Wrap example code in .. code-block:: python directive so Sphinx renders with syntax highlighting."
Suggested change
| - Wrap example code in `.. code-block:: python` directive so Sphinx render with syntax highlighting. | |
| - Wrap example code in `.. code-block:: python` directive so Sphinx renders with syntax highlighting. |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
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
1 participant
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.
Fixes the daily CI failure caused by three independent JAX 0.11.0 internal API breaks, and renames the project instruction file. Compatible with all jax ≥ 0.8.0 via capability detection (no hard-coded version checks).
JAX 0.11 fixes
ClosedJaxprmerged intoJaxpr— the unified constructor dropped thejaxpr=keyword;ModuleInfo.add_jaxpr_outsnow passesjaxpr/constspositionally.jax.core.new_jaxpr_eqnremoved — annotate the older-JAX import fallback withtype: ignoreso mypy passes on new JAX.num_consts/num_carry/linear/_split_transposeparams were removed; the input/output splits now live inft_in/ft_outflattrees. Added two capability-detecting shims in_compatible_imports(scan_num_consts_carry,scan_params_add_ys), routed allnum_consts/num_carryreads through them, and rebuildft_outinadd_scan_ys.while/condparams are unchanged in 0.11 and untouched.CI
0.10.0to the jax-version matrix (now covers0.8.0,0.9.0,0.10.0, latest — old encoding + flattree).Docs / meta
CLAUDE.mdto a trackedAGENTS.md(contents unchanged, then working-agreement rules refined: spec location →docs/specs).Verification
mypy braintrace: clean (59 files).🤖 Generated with Claude Code
Summary by Sourcery
Make braintrace compatible with JAX 0.11's updated jaxpr and scan flattree APIs while retaining support for older JAX versions, expand CI JAX coverage, and introduce a tracked AGENTS.md project guidance document.
Bug Fixes:
Enhancements:
CI:
Documentation: