Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Sorry @r-fedorov, your pull request is larger than the review limit of 300000 diff characters
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (11)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR adds MLX-compatible spherical and Cartesian tensor helpers, tutorial notebooks and compatibility documentation, intermediate and output irrep filtering for ChangesReduced tensor-product filtering
Custom-kernel routing
Estimated code review effort: 4 (Complex) | ~75 minutes 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)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tutorials/other_helpers.py (1)
140-152: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the dead, buggy first
plotdefinition (shadowed by Lines 210-243).This
plotis overridden by the later definition of the same name, so it is unreachable. It also contains a bug: at Line 149d['surfacecolor'].abs()is called on a NumPy array (tracerunsnp.asarray), which has no.abs()method and would raiseAttributeError. Delete this duplicate to avoid confusion.🤖 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 `@tutorials/other_helpers.py` around lines 140 - 152, Remove the earlier duplicate plot function definition, including its trace-building and figure-display logic. Keep the later plot definition around Lines 210–243 as the sole implementation, since the earlier version is shadowed and contains the invalid d['surfacecolor'].abs() call.tests/test_upstream_o3_reduce_tensor.py (1)
107-128: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the new filter validation paths.
Only the combined happy-path (
filter_ir_mid+filter_ir_outboth valid) is tested. The newValueErrorbranches inReducedTensorProducts.__init__for invalidfilter_ir_mid/filter_ir_outcontent, and the case of using each filter independently, are untested.♻️ Suggested additional test
def test_upstream_reduced_tensor_rejects_invalid_filters() -> None: irreps = o3.Irreps.spherical_harmonics(2) with pytest.raises(ValueError, match="filter_ir_mid"): o3.ReducedTensorProducts("ij=ji", i=irreps, filter_ir_mid=[object()]) with pytest.raises(ValueError, match="filter_ir_out"): o3.ReducedTensorProducts("ij=ji", i=irreps, filter_ir_out=[object()])🤖 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 `@tests/test_upstream_o3_reduce_tensor.py` around lines 107 - 128, Add tests alongside test_upstream_reduced_tensor_supports_intermediate_and_output_filters for ReducedTensorProducts filter validation: assert ValueError with matching “filter_ir_mid” and “filter_ir_out” messages when each filter receives invalid content, and add separate cases exercising filter_ir_mid alone and filter_ir_out alone with valid filters.
🤖 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 `@tests/test_documentation_contract.py`:
- Around line 181-185: Correct the contract entry in
tests/test_documentation_contract.py by keeping the filter_ir_mid claim while
pointing its source to an existing tracked documentation file containing the
exact excerpt, or restore docs/guide/tutorial_compatibility.md with that claim.
Ensure the referenced path and excerpt both resolve without changing the
contract text or associated test symbol.
In `@tutorials/data_types.ipynb`:
- Around line 629-630: Update the matrix construction near CartesianTensor to
transpose M with axes (1, 0) instead of (0, 1), so M + M.transpose(1, 0)
actually symmetrizes the matrix while preserving the existing formula and
conversion flow.
In `@tutorials/other_helpers.py`:
- Around line 201-206: Update random_point_signal to use MLX APIs throughout:
generate points with mx.random.normal(shape=(n, 3)), normalize them with
mx.linalg.norm(points, axis=-1, keepdims=True), and construct the result via
SphericalTensor.from_geometry(points, lmax) instead of the foreign
io.SphericalTensor/from_geometry_adjusted call.
---
Nitpick comments:
In `@tests/test_upstream_o3_reduce_tensor.py`:
- Around line 107-128: Add tests alongside
test_upstream_reduced_tensor_supports_intermediate_and_output_filters for
ReducedTensorProducts filter validation: assert ValueError with matching
“filter_ir_mid” and “filter_ir_out” messages when each filter receives invalid
content, and add separate cases exercising filter_ir_mid alone and filter_ir_out
alone with valid filters.
In `@tutorials/other_helpers.py`:
- Around line 140-152: Remove the earlier duplicate plot function definition,
including its trace-building and figure-display logic. Keep the later plot
definition around Lines 210–243 as the sole implementation, since the earlier
version is shadowed and contains the invalid d['surfacecolor'].abs() call.
🪄 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: c2c1ac1f-be7b-44e5-b489-ede5ad31a625
📒 Files selected for processing (13)
.gitignoredocs/guide/tensor_products.mde3nn_mlx/ops_reduce_tensor.pytests/test_documentation_contract.pytests/test_tutorial_tensor_helpers.pytests/test_upstream_o3_reduce_tensor.pytutorials/README.mdtutorials/__init__.pytutorials/data_types.ipynbtutorials/e3nn_MRS_invariants_tutorial_lecture_pt1.ipynbtutorials/operations_on_spherical_tensors.ipynbtutorials/other_helpers.pytutorials/tensor_helpers.py
| ( | ||
| "docs/guide/tutorial_compatibility.md", | ||
| "`filter_ir_mid` restricts every sequential Clebsch--Gordan coupling, including the final coupling.", | ||
| "tests/test_upstream_o3_reduce_tensor.py::test_upstream_reduced_tensor_supports_intermediate_and_output_filters", | ||
| ), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Confirm docs/guide/tutorial_compatibility.md exists and contains the exact excerpt.
fd -HI 'tutorial_compatibility.md'
rg -n "filter_ir_mid.*restricts every sequential Clebsch" docs 2>/dev/nullRepository: lamalab-org/e3nn_mlx
Length of output: 158
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate files =="
fd -HI 'tutorial_compatibility.md|tensor_products.md|documentation_contract.py' .
echo
echo "== relevant contract entry and assertions =="
FILE="$(fd -HI 'test_documentation_contract.py' . | head -n 1)"
if [ -n "${FILE:-}" ]; then
wc -l "$FILE"
sed -n '160,195p;400,430p' "$FILE" | cat -n
fi
echo
echo "== exact excerpt search (all tracked text) =="
python3 - <<'PY'
from pathlib import Path
needle = '`filter_ir_mid` restricts every sequential Clebsch--Gordan coupling, including the final coupling.'
hits=[]
for p in Path('.').rglob('*'):
if p.is_file() and not any(part.startswith('.') for part in p.parts):
try:
text=p.read_text(encoding='utf-8')
except Exception:
continue
if needle in text:
lines=[n for n,line in enumerate(text.splitlines(),1) if needle in line]
hits.append((str(p), lines[:3]))
if hits:
for path, lines in hits:
print(f"{path}: {lines}")
else:
print("NO HITS")
PY
echo
echo "== docs/guide files mentioning filter_ir_mid or Clebsch =="
rg -n "filter_ir_mid|Clebsch" docs 2>/dev/null || trueRepository: lamalab-org/e3nn_mlx
Length of output: 4436
Add or correct the documented source for this contract entry.
"docs/guide/tutorial_compatibility.md" is not present, and the exact excerpt does not exist anywhere under tracked files. Keep the contract but either restore the intended file or update the target to existing documentation containing this claim.
🤖 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 `@tests/test_documentation_contract.py` around lines 181 - 185, Correct the
contract entry in tests/test_documentation_contract.py by keeping the
filter_ir_mid claim while pointing its source to an existing tracked
documentation file containing the exact excerpt, or restore
docs/guide/tutorial_compatibility.md with that claim. Ensure the referenced path
and excerpt both resolve without changing the contract text or associated test
symbol.
| "M = mx.random.normal(shape=(3,3))\n", | ||
| "matrix = CartesianTensor(M + M.transpose(0, 1), formula='ij=ji').to_irrep_tensor()\n", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
M.transpose(0, 1) is a no-op, so this isn't symmetrizing.
transpose(0, 1) is the identity permutation on a 2D array, making M + M.transpose(0, 1) equal to 2*M rather than M + Mᵀ. The intended transpose is transpose(1, 0) (or M.T), consistent with the rank-3/rank-4 cells (Lines 692-694, 731-733). The printed Rs happens to stay correct because formula='ij=ji' projects onto the symmetric irreps regardless, but the input construction is misleading.
🐛 Proposed fix
-matrix = CartesianTensor(M + M.transpose(0, 1), formula='ij=ji').to_irrep_tensor()
+matrix = CartesianTensor(M + M.transpose(1, 0), formula='ij=ji').to_irrep_tensor()📝 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.
| "M = mx.random.normal(shape=(3,3))\n", | |
| "matrix = CartesianTensor(M + M.transpose(0, 1), formula='ij=ji').to_irrep_tensor()\n", | |
| "M = mx.random.normal(shape=(3,3))\n", | |
| "matrix = CartesianTensor(M + M.transpose(1, 0), formula='ij=ji').to_irrep_tensor()\n", |
🤖 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 `@tutorials/data_types.ipynb` around lines 629 - 630, Update the matrix
construction near CartesianTensor to transpose M with axes (1, 0) instead of (0,
1), so M + M.transpose(1, 0) actually symmetrizes the matrix while preserving
the existing formula and conversion flow.
| def random_point_signal(lmax, N=5, r_min=1e-1, std=1.): | ||
| n = N | ||
| points = mx.norm.random(n, 3) * std / mx.sqrt(3.) | ||
| points /= points.norm(2, -1, keepdim=True) | ||
| # select = (points.norm(2, -1) > r_min).nonzero() | ||
| return io.SphericalTensor(lmax, p_val=1, p_arg=-1).from_geometry_adjusted(points) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
random_point_signal references several nonexistent/foreign APIs and will fail if called.
- Line 203
mx.norm.random(n, 3)is not an MLX API (MLX usesmx.random.normal(shape=(n, 3))). - Line 204
points.norm(2, -1, keepdim=True)is Torch tensor syntax, not MLX (mx.linalg.norm(points, axis=-1, keepdims=True)). - Line 206
io.SphericalTensor(...).from_geometry_adjusted(...)—iohere isase.io; this is a leftover e3nn reference. The MLX helper isSphericalTensor.from_geometry(points, lmax).
This function cannot run as written.
🐛 Proposed fix
-def random_point_signal(lmax, N=5, r_min=1e-1, std=1.):
- n = N
- points = mx.norm.random(n, 3) * std / mx.sqrt(3.)
- points /= points.norm(2, -1, keepdim=True)
-# select = (points.norm(2, -1) > r_min).nonzero()
- return io.SphericalTensor(lmax, p_val=1, p_arg=-1).from_geometry_adjusted(points)
+def random_point_signal(lmax, N=5, r_min=1e-1, std=1.0):
+ points = mx.random.normal(shape=(N, 3)) * std / math.sqrt(3.0)
+ points = points / mx.linalg.norm(points, axis=-1, keepdims=True)
+ return SphericalTensor.from_geometry(points, lmax)📝 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.
| def random_point_signal(lmax, N=5, r_min=1e-1, std=1.): | |
| n = N | |
| points = mx.norm.random(n, 3) * std / mx.sqrt(3.) | |
| points /= points.norm(2, -1, keepdim=True) | |
| # select = (points.norm(2, -1) > r_min).nonzero() | |
| return io.SphericalTensor(lmax, p_val=1, p_arg=-1).from_geometry_adjusted(points) | |
| def random_point_signal(lmax, N=5, r_min=1e-1, std=1.0): | |
| points = mx.random.normal(shape=(N, 3)) * std / math.sqrt(3.0) | |
| points = points / mx.linalg.norm(points, axis=-1, keepdims=True) | |
| return SphericalTensor.from_geometry(points, lmax) |
🤖 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 `@tutorials/other_helpers.py` around lines 201 - 206, Update
random_point_signal to use MLX APIs throughout: generate points with
mx.random.normal(shape=(n, 3)), normalize them with mx.linalg.norm(points,
axis=-1, keepdims=True), and construct the result via
SphericalTensor.from_geometry(points, lmax) instead of the foreign
io.SphericalTensor/from_geometry_adjusted call.
1 participant
Summary by CodeRabbit
use_custom_kerneltoggle for v2.1.0.6 point-cloud models, propagated through message passing and convolutions.