Conversation
Reviewer's GuideAdds JVP-safe scatter_sum fallback and extensive documentation and tooling for the public docs site, including Sphinx configuration, GitHub Pages deployment, and executable documentation contracts in tests. Flow diagram for scatter_sum JVP-safe fallback selectionflowchart TD
start["call scatter_sum(source, index, dim_size, use_custom_kernel, jvp_safe)"]
validate_dim["validate dim_size and infer when None"]
check_jvp["jvp_safe?"]
check_kernel["use_custom_kernel?"]
error_kernel["raise ValueError(jvp_safe_requires_general_path)"]
call_fixed["_fixed_index_scatter_sum(source, index, dim_size)"]
check_metal["mlx_metal_available and use_custom_kernel"]
custom_kernel["use _scatter_sum_custom_kernel(source, index, dim_size)"]
general_path["use output.at[index].add(source)"]
start --> validate_dim --> check_jvp
check_jvp -->|true| check_kernel
check_jvp -->|false| check_metal
check_kernel -->|true| error_kernel
check_kernel -->|false| call_fixed
check_metal -->|true| custom_kernel
check_metal -->|false| general_path
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
📝 WalkthroughWalkthroughThe change adds a Sphinx documentation site with local build tooling and GitHub Pages deployment, expands guides and API references, introduces JVP-safe fixed-index scatter reduction, and adds executable documentation and runtime contract tests. ChangesDocumentation and JVP-safe scatter support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant scatter_sum
participant _fixed_index_scatter_sum
participant MLX
Caller->>scatter_sum: pass jvp_safe=True
scatter_sum->>_fixed_index_scatter_sum: dispatch fixed-index reduction
_fixed_index_scatter_sum->>MLX: compute sorted prefix sums
MLX-->>Caller: return reduced values and JVP-compatible results
Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Checkov (3.3.8).github/workflows/docs.ymlTraceback (most recent call last): Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 1 issue, and left some high level feedback:
- The
_fixed_index_scatter_sumfallback performsnp.searchsortedovernp.arange(dim_size), which is O(dim_size) even for very sparse indices; consider constraining this to the actual min/max occupied rows or documenting the runtime/memory tradeoff for largedim_size. _fixed_index_scatter_sumimports NumPy inside the function on every call; moving this import to module scope would avoid repeated import overhead and make the function cheaper in hot paths that rely onjvp_safe=True.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `_fixed_index_scatter_sum` fallback performs `np.searchsorted` over `np.arange(dim_size)`, which is O(dim_size) even for very sparse indices; consider constraining this to the actual min/max occupied rows or documenting the runtime/memory tradeoff for large `dim_size`.
- `_fixed_index_scatter_sum` imports NumPy inside the function on every call; moving this import to module scope would avoid repeated import overhead and make the function cheaper in hot paths that rely on `jvp_safe=True`.
## Individual Comments
### Comment 1
<location path="docs/guide/installation.md" line_range="32-33" />
<code_context>
+python -m pytest
+```
+
+MLX tests that require a working runtime can otherwise skip. CI prevents an
+accidental pass by setting `E3NN_MLX_REQUIRE_RUNTIME=1`.
+
</code_context>
<issue_to_address>
**suggestion (typo):** Consider rephrasing this sentence for clearer grammar.
The phrase "can otherwise skip" is a bit unclear, as it sounds like the tests perform the action. Consider changing it to "can otherwise be skipped" or "can otherwise be omitted" for more natural phrasing.
```suggestion
MLX tests that require a working runtime can otherwise be skipped. CI prevents an
accidental pass by setting `E3NN_MLX_REQUIRE_RUNTIME=1`.
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+32
to
+33
| MLX tests that require a working runtime can otherwise skip. CI prevents an | ||
| accidental pass by setting `E3NN_MLX_REQUIRE_RUNTIME=1`. |
There was a problem hiding this comment.
suggestion (typo): Consider rephrasing this sentence for clearer grammar.
The phrase "can otherwise skip" is a bit unclear, as it sounds like the tests perform the action. Consider changing it to "can otherwise be skipped" or "can otherwise be omitted" for more natural phrasing.
Suggested change
| MLX tests that require a working runtime can otherwise skip. CI prevents an | |
| accidental pass by setting `E3NN_MLX_REQUIRE_RUNTIME=1`. | |
| MLX tests that require a working runtime can otherwise be skipped. CI prevents an | |
| accidental pass by setting `E3NN_MLX_REQUIRE_RUNTIME=1`. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/conf.py`:
- Line 10: Add a targeted Ruff suppression for the required global copyright
setting in docs/conf.py, preferably an inline # noqa: A001 on the copyright
assignment; do not alter the Sphinx variable or broaden lint ignores.
In `@docs/development/deployment.md`:
- Around line 20-23: Revise the deployment documentation to avoid claiming the
pull-request documentation build is required unless repository branch protection
or ruleset configuration is explicitly documented. Prefer stating that PRs build
the site, or add the corresponding required-check configuration details.
In `@docs/examples/point_models.md`:
- Around line 40-48: Update the forward helper around model.forward_with_edges
to derive num_graphs from the supplied graph_batch instead of hard-coding 1,
ensuring batched inputs produce the correct pooling output size; alternatively,
explicitly enforce and document a single-graph graph_batch precondition.
In `@docs/guide/installation.md`:
- Around line 37-41: Update the scatter guidance in the installation
documentation to state that forward-mode JVPs require `jvp_safe=True` with
eager, fixed indices; do not imply that `use_custom_kernel=False` alone provides
JVP support. Preserve the existing general-operation guidance and distinguish it
from the JVP-safe execution path.
In `@e3nn_mlx/graph.py`:
- Around line 83-87: Update the row-label construction in the graph grouping
logic around order_host and the starts_host/ends_host searchsorted calls to use
a wide signed dtype instead of host_index.dtype, preventing valid narrow
unsigned indices from wrapping when dim_size exceeds their range. Add a
regression test covering narrow unsigned indices with dim_size=257 and verify
row 256 remains in its own bucket.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 93dda358-a21e-4716-a10a-61a91f1d728d
📒 Files selected for processing (34)
.github/workflows/docs.yml.gitignoreREADME.mddocs/HIGH_LEVEL_API.mddocs/Makefiledocs/_static/.gitkeepdocs/_templates/.gitkeepdocs/api/graph.rstdocs/api/index.mddocs/api/math.rstdocs/api/models.rstdocs/api/nn.rstdocs/api/o3.rstdocs/api/typed.rstdocs/conf.pydocs/development/deployment.mddocs/development/index.mddocs/examples/convolution.mddocs/examples/getting_started.mddocs/examples/index.mddocs/examples/point_models.mddocs/guide/equivariance.mddocs/guide/index.mddocs/guide/installation.mddocs/guide/irreps.mddocs/guide/migration.mddocs/guide/performance.mddocs/guide/tensor_products.mddocs/index.mde3nn_mlx/graph.pyevals/KERNEL_EVALUATION.mdpyproject.tomltests/test_documentation_contract.pytests/test_graph_radial.py
|
|
||
| project = "e3nn-mlx" | ||
| author = "e3nn-mlx contributors" | ||
| copyright = "2026, lamalab-org" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Files matching conf.py/build/lint config:\n'
git ls-files | rg '(^|/)(docs/conf\.py|pyproject\.toml|ruff\.toml|\.ruff\.toml|tox\.ini|pre-commit\.config\.yaml|noxfile\.py)$' || true
printf '\ndocs/conf.py excerpt:\n'
if [ -f docs/conf.py ]; then
sed -n '1,80p' docs/conf.py | cat -n
fi
printf '\nruff-related occurrences:\n'
rg -n "ruff|per-file-ignores|rule-select|ignore|noqa" pyproject.toml ruff.toml .ruff.toml pre-commit.config.yaml tox.ini docs/conf.py 2>/dev/null || true
printf '\nSearch for A001/copyright variable in tracked files:\n'
rg -n "A001|builtins\.copyright|noqa|copyright" -S .Repository: lamalab-org/e3nn_mlx
Length of output: 1969
Supress Ruff A001 for Sphinx’s required setting.
Sphinx’s docs/conf.py requires the global copyright variable, but Ruff will see it as shadowing the built-in. Add a targeted # noqa: A001 or per-file ignore for docs/conf.py.
Proposed fix
-copyright = "2026, lamalab-org"
+copyright = "2026, lamalab-org" # noqa: A001📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| copyright = "2026, lamalab-org" | |
| copyright = "2026, lamalab-org" # noqa: A001 |
🧰 Tools
🪛 Ruff (0.15.21)
[error] 10-10: Variable copyright is shadowing a Python builtin
(A001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/conf.py` at line 10, Add a targeted Ruff suppression for the required
global copyright setting in docs/conf.py, preferably an inline # noqa: A001 on
the copyright assignment; do not alter the Sphinx variable or broaden lint
ignores.
Source: Linters/SAST tools
Comment on lines
+20
to
+23
| The repository includes `.github/workflows/docs.yml`. Pull requests build the | ||
| site as a required-quality check but do not deploy it. A push to `main` that | ||
| changes documentation, package source, or documentation configuration builds | ||
| and publishes the Pages artifact. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Do not imply that the PR check is required without repository enforcement.
The workflow can build documentation for pull requests, but GitHub only makes that check mandatory when branch protection or a ruleset marks it as required. Either reword this as “PRs build the site” or document the required-check configuration.
🧰 Tools
🪛 LanguageTool
[uncategorized] ~20-~20: The official name of this software platform is spelled with a capital “H”.
Context: ...h GitHub Pages The repository includes .github/workflows/docs.yml. Pull requests buil...
(GITHUB)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/development/deployment.md` around lines 20 - 23, Revise the deployment
documentation to avoid claiming the pull-request documentation build is required
unless repository branch protection or ruleset configuration is explicitly
documented. Prefer stating that PRs build the site, or add the corresponding
required-check configuration details.
Comment on lines
+40
to
+48
| def forward(pos, features, graph_batch, src, dst): | ||
| return model.forward_with_edges( | ||
| pos, | ||
| features, | ||
| graph_batch, | ||
| src, | ||
| dst, | ||
| num_graphs=1, | ||
| ).array |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Derive num_graphs from the supplied batch.
SimpleNetwork.forward_with_edges uses num_graphs as the output size for node pooling. Hard-coding 1 breaks valid batched inputs whose graph_batch contains more than one graph. Compute the count outside the compiled function for the fixed batch, or explicitly document and assert a single-graph precondition. (raw.githubusercontent.com)
Proposed fix
+num_graphs = (
+ 0 if data["batch"].shape[0] == 0
+ else int(mx.max(data["batch"]).item()) + 1
+)
+
def forward(pos, features, graph_batch, src, dst):
return model.forward_with_edges(
pos,
features,
graph_batch,
src,
dst,
- num_graphs=1,
+ num_graphs=num_graphs,
).array🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/examples/point_models.md` around lines 40 - 48, Update the forward
helper around model.forward_with_edges to derive num_graphs from the supplied
graph_batch instead of hard-coding 1, ensuring batched inputs produce the
correct pooling output size; alternatively, explicitly enforce and document a
single-graph graph_batch precondition.
Comment on lines
+37
to
+41
| Generated kernels are enabled by default on compatible Apple-silicon inputs. | ||
| Every public operation retains a general MLX path. For spherical harmonics and | ||
| scatter, pass `use_custom_kernel=False`; for tensor products, construct the | ||
| module with `use_custom_kernel=False` or call `differentiable_arrays` when a | ||
| forward-mode transform is needed. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Document jvp_safe=True for scatter JVPs.
scatter_sum(..., use_custom_kernel=False) still uses the indexed-add path, which has no JVP rule. Forward-mode differentiation requires jvp_safe=True with eager, fixed indices; update this guidance to distinguish general execution from the JVP-safe path. (raw.githubusercontent.com)
Proposed wording
-For spherical harmonics and scatter, pass `use_custom_kernel=False`; for tensor
-products, construct the module with `use_custom_kernel=False` or call
-`differentiable_arrays` when a forward-mode transform is needed.
+For spherical harmonics, pass `use_custom_kernel=False`. For `scatter_sum`,
+use `jvp_safe=True` with eager, fixed indices when forward-mode differentiation
+is needed; this path requires `use_custom_kernel=False`. For tensor products,
+construct the module with `use_custom_kernel=False` or call
+`differentiable_arrays` when a forward-mode transform is needed.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| Generated kernels are enabled by default on compatible Apple-silicon inputs. | |
| Every public operation retains a general MLX path. For spherical harmonics and | |
| scatter, pass `use_custom_kernel=False`; for tensor products, construct the | |
| module with `use_custom_kernel=False` or call `differentiable_arrays` when a | |
| forward-mode transform is needed. | |
| Generated kernels are enabled by default on compatible Apple-silicon inputs. | |
| Every public operation retains a general MLX path. For spherical harmonics, pass `use_custom_kernel=False`. For `scatter_sum`, | |
| use `jvp_safe=True` with eager, fixed indices when forward-mode differentiation | |
| is needed; this path requires `use_custom_kernel=False`. For tensor products, | |
| construct the module with `use_custom_kernel=False` or call `differentiable_arrays` | |
| when a forward-mode transform is needed. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/guide/installation.md` around lines 37 - 41, Update the scatter guidance
in the installation documentation to state that forward-mode JVPs require
`jvp_safe=True` with eager, fixed indices; do not imply that
`use_custom_kernel=False` alone provides JVP support. Preserve the existing
general-operation guidance and distinguish it from the JVP-safe execution path.
Comment on lines
+83
to
+87
| order_host = np.argsort(host_index, kind="stable") | ||
| sorted_index = host_index[order_host] | ||
| rows = np.arange(dim_size, dtype=host_index.dtype) | ||
| starts_host = np.searchsorted(sorted_index, rows, side="left") | ||
| ends_host = np.searchsorted(sorted_index, rows, side="right") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
python - <<'PY'
import numpy as np
index = np.array([0, 1], dtype=np.uint8)
rows = np.arange(257, dtype=index.dtype)
assert rows[256] == 0 # Demonstrates the erroneous wrapped lookup label.
PYRepository: lamalab-org/e3nn_mlx
Length of output: 158
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file =="
git ls-files | rg '(^|/)e3nn_mlx/graph\.py$' || true
echo "== file context =="
if [ -f e3nn_mlx/graph.py ]; then
wc -l e3nn_mlx/graph.py
sed -n '1,140p' e3nn_mlx/graph.py | cat -n
fi
echo "== relevant occurrences =="
rg -n "def .*|host_index|order_host|sorted_index|rows = |searchsorted|dim_size" e3nn_mlx/graph.py || true
echo "== behavioral probe of NumPy semantics with uint8 rows =="
python3 - <<'PY'
import numpy as np
for dtype in [np.uint8, np.uint16, np.int8, np.int16, np.int64]:
for rows_max in [255, 256, 257, 766, 768]:
rows = np.arange(rows_max, dtype=dtype)
wrapped = rows[-1] if int(rows_max) < 2**64 or dtype in (np.int8, np.int16) else None
print(dtype, rows_max, wrapped)
PYRepository: lamalab-org/e3nn_mlx
Length of output: 7772
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repo-wide occurrences of jvp_safe scatter =="
rg -n "jvp_safe|\bscatter_sum\(" . --glob '!*.log' || true
echo "== tests mentioning fixed_index or scatter_sum =="
rgitz -n "fixed_index_scatter_sum|_fixed_index_scatter_sum|scatter_sum|jvp_safe|uint8.*scatter|dim_size.*256|dim_size.*257" --glob '!*.log' . || true
echo "== file size and test files =="
wc -l e3nn_mlx/graph.py
git ls-files | rg '(^|/)(test_|tests|.*test.*\.py$)' || trueRepository: lamalab-org/e3nn_mlx
Length of output: 5897
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== tests/test_graph_radial.py around scatter_sum =="
sed -n '100,185p' tests/test_graph_radial.py | cat -n
echo "== tests/test_documentation_contract.py around scatter_sum =="
sed -n '600,660p' tests/test_documentation_contract.py | cat -n
echo "== tests/test_metal_kernels.py around scatter_sum =="
sed -n '95,125p' tests/test_metal_kernels.py | cat -nRepository: lamalab-org/e3nn_mlx
Length of output: 8511
Avoid wrapping row labels for narrow index dtypes.
rows inherits host_index.dtype, so valid uint8 indices can wrap at dim_size=257 and send the 256th row into bucket 0. Use a wide signed dtype for the lookup labels and add a regression test for narrow unsigned indices.
Proposed fix
- rows = np.arange(dim_size, dtype=host_index.dtype)
+ rows = np.arange(dim_size, dtype=np.int64)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| order_host = np.argsort(host_index, kind="stable") | |
| sorted_index = host_index[order_host] | |
| rows = np.arange(dim_size, dtype=host_index.dtype) | |
| starts_host = np.searchsorted(sorted_index, rows, side="left") | |
| ends_host = np.searchsorted(sorted_index, rows, side="right") | |
| order_host = np.argsort(host_index, kind="stable") | |
| sorted_index = host_index[order_host] | |
| rows = np.arange(dim_size, dtype=np.int64) | |
| starts_host = np.searchsorted(sorted_index, rows, side="left") | |
| ends_host = np.searchsorted(sorted_index, rows, side="right") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@e3nn_mlx/graph.py` around lines 83 - 87, Update the row-label construction in
the graph grouping logic around order_host and the starts_host/ends_host
searchsorted calls to use a wide signed dtype instead of host_index.dtype,
preventing valid narrow unsigned indices from wrapping when dim_size exceeds
their range. Add a regression test covering narrow unsigned indices with
dim_size=257 and verify row 256 remains in its own bucket.
Merged
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.
Summary by Sourcery
Document and test the JVP-safe scatter path while adding a full Sphinx documentation site and associated CI workflow.
New Features:
Bug Fixes:
Enhancements:
CI:
Tests:
Summary by CodeRabbit
New Features
Documentation
Bug Fixes
Tests