Conversation
![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 (13)
🚧 Files skipped from review as they are similar to previous changes (9)
📝 WalkthroughWalkthroughThis change corrects tensor-product and ChangesTensor semantics and validation
Three-way benchmark runner
Randomized parity qualification
Compatibility and evaluation documentation
Estimated code review effort: 5 (Critical) | ~120 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 |
Reviewer's GuideRefactors the evaluation and benchmarking suite into a compact three‑way Torch‑CPU vs MLX vs MLX‑kernel benchmark, tightens tensor product and Linear semantics and validation, adds randomized cross‑framework numerical parity harnesses, and augments Metal kernel dispatch logic and tests, while simplifying plots, workloads, and documentation accordingly. Sequence diagram for MLX TensorProduct Metal dispatch vs general pathsequenceDiagram
participant Caller
participant TP as TensorProduct
participant MDK as TensorProduct__metal_dispatch_kind
participant TM as TensorProduct__try_metal
participant Metal as _metal_operation
participant Num as _numerical_tensor_product
Caller->>TP: _apply_arrays(left_array, right_array, weight)
TP->>TP: _get_weight(weight)
alt use_custom_kernel enabled
TP->>TM: _try_metal(left_array, right_array, weight)
TM->>MDK: _metal_dispatch_kind(left_array, right_array, weight)
MDK-->>TM: kind or None
alt kind is not None
TM->>Metal: _metal_operation(left_array, right_array, metal_weight)
Metal-->>TP: output_array
else kind is None
TM-->>TP: None
TP->>Num: _numerical_tensor_product(left, right, weights)
Num-->>TP: output_array
end
else use_custom_kernel disabled
TP->>Num: _numerical_tensor_product(left, right, weights)
Num-->>TP: output_array
end
TP-->>Caller: output_array
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![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 (6)
tests/test_metal_kernels.py (1)
214-221: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMatching on an MLX-internal error message is brittle.
"Not implemented for CustomKernel"is MLX's own wording for the missing JVP rule, not this project's contract, so an upstream rewording breaks the test for reasons unrelated toe3nn_mlx. Consider asserting only the exception type, or relaxing the pattern to a stable fragment such as"CustomKernel".♻️ Suggested relaxation
- with pytest.raises(ValueError, match="Not implemented for CustomKernel"): + with pytest.raises(ValueError, match="CustomKernel"):🤖 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_metal_kernels.py` around lines 214 - 221, Relax the ValueError assertion in the mx.jvp test around apply so it does not depend on MLX’s exact internal message wording. Assert only ValueError, or match the stable “CustomKernel” fragment while preserving the test’s validation of the unsupported JVP behavior.tests/test_tensor_product_compatibility_qualification.py (1)
278-304: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThis test can pass vacuously.
record_broadcastonly records whenarray is shared_weight. If MLX ever returns a distinct array object fromproduct.weight, the identity check never fires andmaterialized_shapes == []holds regardless of whether the weight was materialized per item — silently losing the performance contract being guarded. Consider matching on shape instead (any broadcast whose trailing dim isweight_numel), and/or asserting the recorder observed at least one call so the hook is proven live.♻️ Suggested strengthening
materialized_shapes = [] + all_broadcasts = [] original_broadcast_to = mx.broadcast_to def record_broadcast(array, shape, *args, **kwargs): - if array is shared_weight: + all_broadcasts.append(tuple(shape)) + if tuple(shape)[-1:] == (product.weight_numel,) and len(shape) > 1: materialized_shapes.append(tuple(shape)) return original_broadcast_to(array, shape, *args, **kwargs) monkeypatch.setattr(mx, "broadcast_to", record_broadcast) output = product(left, right) mx.eval(output.array) assert output.shape == (8, product.irreps_out.dim) assert materialized_shapes == []🤖 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_tensor_product_compatibility_qualification.py` around lines 278 - 304, Strengthen test_shared_tensor_product_weights_are_not_materialized_per_item so the broadcast recorder cannot pass vacuously when product.weight is represented by a different array object. Match relevant broadcasts using the shared weight shape or trailing weight_numel dimension, and assert the recorder observed the expected hook activity before verifying that no per-item materialization occurred.tests/test_upstream_o3_linear_norm.py (1)
156-196: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a sparse-group case to cover the zero-block branch.
With
3x2e + 3x3e + 2x2e→3x2e + 3x3e + 3x2eand default instructions, the2egroup has all four(input, output)pairs present, so the newmx.zeros(...)fallback innn_linear.py(lines 294-302) is never taken. That branch is exactly where a wrong block shape or ordering would silently corrupt the assembled matrix. Passing explicitinstructionsthat omit one pair within the2egroup (e.g.[(0, 0), (1, 1), (2, 2)]) would exercise it alongside this VJP formula.🤖 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_linear_norm.py` around lines 156 - 196, Update test_grouped_linear_external_weight_vjp_matches_blockwise_formula to pass explicit instructions omitting one 2e input/output pair, such as the diagonal pairs (0, 0), (1, 1), and (2, 2), when constructing o3.Linear. Keep the existing VJP-versus-blockwise gradient assertions unchanged so the test exercises the zero-block fallback while preserving the formula coverage.tests/test_randomized_tensor_product_parity_harness.py (1)
77-90: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSelect the perturbed variant by its kernel flag instead of
index == 2.
MLX_VARIANTSentries are(name, compile_left_right, use_custom_kernel); keying off position 2 breaks whenever the variant list is reordered or extended.♻️ Proposed tweak
- for index, (name, _, _) in enumerate(MLX_VARIANTS): + for name, _compile, use_custom_kernel in MLX_VARIANTS: output = reference.copy() - if index == 2: + if use_custom_kernel: output[0, 0] += 0.1 variants.append( { "name": name, "status": "ok", "shape": [1, 2], "output": output.tolist(), - "custom_eligible": index == 2, - "kernel_kind": "scalar_paths" if index == 2 else None, + "custom_eligible": use_custom_kernel, + "kernel_kind": "scalar_paths" if use_custom_kernel else None, } )🤖 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_randomized_tensor_product_parity_harness.py` around lines 77 - 90, Update the variant-generation loop to select the perturbed and custom-kernel variant using each MLX_VARIANTS entry’s use_custom_kernel flag rather than index == 2. Unpack the tuple accordingly and derive both custom_eligible and kernel_kind from that flag, preserving the existing scalar_paths and default behavior.evals/torch_cases.py (2)
144-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRuff flags these
lambdaassignments as E731 errors; convert todef.Same pattern at Lines 144, 215, 307, 339, 385. Ruff also reports the unpacked-but-unused
torchat Lines 270, 318, 350 (RUF059) — use_torchor_, o3 = _imports()there.♻️ Example
- forward = lambda: module(left, right) + def forward(): + return module(left, right)Also applies to: 215-215, 307-307, 339-339, 385-385
🤖 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 `@evals/torch_cases.py` at line 144, Replace the lambda assignments at the `forward` definitions on lines 144, 215, 307, 339, and 385 with local `def` functions preserving each callable’s arguments and behavior. In the affected cases at lines 270, 318, and 350, rename the unused unpacked `torch` binding to `_torch` or discard it with `_` while retaining the `o3` value.Source: Linters/SAST tools
300-300: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGlobal monkeypatch of
model_module.radius_graphleaks across builders.Each builder permanently rebinds the attribute on the imported upstream module with a closure over its
edge_index. It works only becauserun_backend.mainconstructs and benchmarks one case at a time; any future change that builds all tasks up front (or reuses a task after a later build) would silently benchmark the wrong topology. Consider patching within a context manager around the measured calls, or passing precomputed edges via the data dict where the upstream API allows it.Also applies to: 337-337
🤖 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 `@evals/torch_cases.py` at line 300, Replace the permanent assignments to model_module.radius_graph in each builder with scoped patching that is active only during that builder’s measured model calls, then restores the original function afterward. Ensure closures use the current builder’s edge_index and cannot affect later builders or reused tasks.
🤖 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 `@evals/mlx_cases.py`:
- Around line 126-134: The dispatch labels in evals/mlx_cases.py at lines
126-134 and 303-305 must reflect runtime kernel eligibility, not
_USE_CUSTOM_KERNELS alone. Update the spherical-harmonics label near the
dispatch expression to also require the same Metal availability conditions used
by ops_sh.spherical_harmonics, and update the scatter-sum label to require the
scatter kernel’s actual eligibility/availability; otherwise retain the existing
fallback labels.
In `@evals/randomized_core_parity.py`:
- Around line 890-913: Pass the use_custom_kernel value through the TensorSquare
construction in the wrapper branch, ensuring the inherited _try_metal probe and
_metal_kernel_kind metadata use the requested custom-kernel setting rather than
the constructor default.
In `@tests/test_documentation_contract.py`:
- Around line 282-341: Update the DOCUMENTED_SOURCE_CONTRACTS entries in
test_documentation_contract.py so every quoted excerpt exactly matches the
current wording and punctuation in its referenced documentation file. Inspect
the corresponding sections of tensor_products.md, linear.md, performance.md, and
evals/README.md, then replace only the stale literal substrings while preserving
each associated test reference and contract intent.
---
Nitpick comments:
In `@evals/torch_cases.py`:
- Line 144: Replace the lambda assignments at the `forward` definitions on lines
144, 215, 307, 339, and 385 with local `def` functions preserving each
callable’s arguments and behavior. In the affected cases at lines 270, 318, and
350, rename the unused unpacked `torch` binding to `_torch` or discard it with
`_` while retaining the `o3` value.
- Line 300: Replace the permanent assignments to model_module.radius_graph in
each builder with scoped patching that is active only during that builder’s
measured model calls, then restores the original function afterward. Ensure
closures use the current builder’s edge_index and cannot affect later builders
or reused tasks.
In `@tests/test_metal_kernels.py`:
- Around line 214-221: Relax the ValueError assertion in the mx.jvp test around
apply so it does not depend on MLX’s exact internal message wording. Assert only
ValueError, or match the stable “CustomKernel” fragment while preserving the
test’s validation of the unsupported JVP behavior.
In `@tests/test_randomized_tensor_product_parity_harness.py`:
- Around line 77-90: Update the variant-generation loop to select the perturbed
and custom-kernel variant using each MLX_VARIANTS entry’s use_custom_kernel flag
rather than index == 2. Unpack the tuple accordingly and derive both
custom_eligible and kernel_kind from that flag, preserving the existing
scalar_paths and default behavior.
In `@tests/test_tensor_product_compatibility_qualification.py`:
- Around line 278-304: Strengthen
test_shared_tensor_product_weights_are_not_materialized_per_item so the
broadcast recorder cannot pass vacuously when product.weight is represented by a
different array object. Match relevant broadcasts using the shared weight shape
or trailing weight_numel dimension, and assert the recorder observed the
expected hook activity before verifying that no per-item materialization
occurred.
In `@tests/test_upstream_o3_linear_norm.py`:
- Around line 156-196: Update
test_grouped_linear_external_weight_vjp_matches_blockwise_formula to pass
explicit instructions omitting one 2e input/output pair, such as the diagonal
pairs (0, 0), (1, 1), and (2, 2), when constructing o3.Linear. Keep the existing
VJP-versus-blockwise gradient assertions unchanged so the test exercises the
zero-block fallback while preserving the formula coverage.
🪄 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: e4e964ac-cb61-4a61-9528-77e89bf0330d
📒 Files selected for processing (39)
CHANGELOG.mdREADME.mddocs/COMPATIBILITY.mddocs/HIGH_LEVEL_API.mddocs/UPSTREAM_O3_TEST_MATRIX.mddocs/guide/index.mddocs/guide/linear.mddocs/guide/migration.mddocs/guide/performance.mddocs/guide/tensor_products.mde3nn_core/instructions.pye3nn_mlx/nn_linear.pye3nn_mlx/ops_tp.pyevals/KERNEL_EVALUATION.mdevals/OPTIMIZATION_REPORT.mdevals/RANDOMIZED_OPERATION_PARITY.mdevals/RANDOMIZED_PARITY.mdevals/README.mdevals/common.pyevals/mlx_cases.pyevals/plot_results.pyevals/randomized_core_parity.pyevals/randomized_operation_parity.pyevals/randomized_tensor_product_parity.pyevals/run.pyevals/run_backend.pyevals/torch_cases.pyevals/workloads.pytests/test_documentation_contract.pytests/test_evals.pytests/test_metal_kernels.pytests/test_p0_equivariance.pytests/test_randomized_core_parity_harness.pytests/test_randomized_operation_parity_harness.pytests/test_randomized_tensor_product_parity_harness.pytests/test_tensor_product.pytests/test_tensor_product_compatibility_qualification.pytests/test_tensor_product_correctness.pytests/test_upstream_o3_linear_norm.py
💤 Files with no reviewable changes (2)
- evals/KERNEL_EVALUATION.md
- evals/OPTIMIZATION_REPORT.md
Comment on lines
+282
to
+341
| ( | ||
| "docs/guide/tensor_products.md", | ||
| "Shared weights remain one-dimensional during execution.", | ||
| "tests/test_tensor_product_compatibility_qualification.py::test_shared_tensor_product_weights_are_not_materialized_per_item", | ||
| ), | ||
| ( | ||
| "docs/guide/tensor_products.md", | ||
| "Scalar-path kernel selection considers both the static contraction size and the batch size.", | ||
| "tests/test_metal_kernels.py::test_dense_scalar_tensor_product_dispatch_uses_work_guard", | ||
| ), | ||
| ( | ||
| "docs/guide/tensor_products.md", | ||
| "Both paths implement the same contraction and normalization conventions.", | ||
| "tests/test_tensor_product_correctness.py::test_optimized_full_tensor_product_matches_fallback_outputs_and_gradients", | ||
| ), | ||
| ( | ||
| "docs/guide/linear.md", | ||
| "Without explicit instructions, `Linear` enumerates compatible paths in input-major, then output-index order, matching upstream e3nn.", | ||
| "tests/test_upstream_o3_linear_norm.py::test_linear_default_weight_paths_follow_upstream_input_major_order", | ||
| ), | ||
| ( | ||
| "docs/guide/linear.md", | ||
| "The grouped execution path preserves derivatives with respect to external weights, including repeated input and output irrep blocks.", | ||
| "tests/test_upstream_o3_linear_norm.py::test_grouped_linear_external_weight_vjp_matches_blockwise_formula", | ||
| ), | ||
| ( | ||
| "docs/COMPATIBILITY.md", | ||
| "Mode `uvw` always requires weights", | ||
| "tests/test_tensor_product_compatibility_qualification.py::test_constructor_rejects_upstream_invalid_uvw_and_static_configurations", | ||
| ), | ||
| ( | ||
| "docs/guide/equivariance.md", | ||
| "actual = layer(x @ D_in.T)", | ||
| "tests/test_documentation_contract.py::test_documented_linear_and_gate_examples_are_equivariant", | ||
| ), | ||
| ( | ||
| "docs/guide/performance.md", | ||
| "The repository's `evals/` harness does this consistently across MLX, PyTorch CPU, and PyTorch MPS.", | ||
| "The repository's `evals/` harness does this consistently across Torch CPU, general MLX, and kernel-enabled MLX.", | ||
| "tests/test_documentation_contract.py::test_documentation_metadata_workflow_and_timing_contracts", | ||
| ), | ||
| ( | ||
| "docs/guide/performance.md", | ||
| "Disabling a kernel changes execution, not the mathematical operation.", | ||
| "tests/test_metal_kernels.py::test_spherical_harmonics_kernel_matches_general_forward_gradient_and_hessian", | ||
| ), | ||
| ( | ||
| "docs/guide/performance.md", | ||
| "Dense shared-weight tensor products deliberately keep their parameter vector compact.", | ||
| "tests/test_tensor_product_compatibility_qualification.py::test_shared_tensor_product_weights_are_not_materialized_per_item", | ||
| ), | ||
| ( | ||
| "docs/guide/performance.md", | ||
| "The shipped `gate_points_2102` model currently has no model-level kernel toggle", | ||
| "tests/test_evals.py::test_model_benchmarks_report_kernel_scope", | ||
| ), | ||
| ( | ||
| "evals/README.md", | ||
| "Every result row also records the selected execution path.", | ||
| "tests/test_evals.py::test_mlx_benchmark_reports_actual_tensor_product_dispatch", | ||
| ), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n 'Shared weights remain one-dimensional during execution.' docs/guide/tensor_products.md
rg -n 'Scalar-path kernel selection considers both the static contraction size and the batch size.' docs/guide/tensor_products.md
rg -n 'Both paths implement the same contraction and normalization conventions.' docs/guide/tensor_products.md
rg -n 'Without explicit instructions, .Linear. enumerates compatible paths in input-major' docs/guide/linear.md
rg -n 'The grouped execution path preserves derivatives with respect to external weights' docs/guide/linear.md
rg -n "The repository's .evals/. harness does this consistently across Torch CPU" docs/guide/performance.md
rg -n 'Dense shared-weight tensor products deliberately keep their parameter vector compact.' docs/guide/performance.md
rg -n 'The shipped .gate_points_2102. model currently has no model-level kernel toggle' docs/guide/performance.md
rg -n 'Every result row also records the selected execution path.' evals/README.mdRepository: lamalab-org/e3nn_mlx
Length of output: 317
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact referenced lines plus surrounding context.
printf 'docs/guide/tensor_products.md lines 65-80 and 108-124:\n'
sed -n '65,80p;108,124p' docs/guide/tensor_products.md
printf '\nSearch exact quoted excerpts (Python literal compare):\n'
python3 - <<'PY'
from pathlib import Path
target_refs = [
("docs/guide/tensor_products.md", "Shared weights remain one-dimensional during execution."),
("docs/guide/tensor_products.md", "Scalar-path kernel selection considers both the static contraction size and the batch size."),
("docs/guide/tensor_products.md", "Both paths implement the same contraction and normalization conventions."),
("docs/guide/linear.md", "Without explicit instructions, `Linear` enumerates compatible paths in input-major, then output-index order, matching upstream e3nn."),
("docs/guide/linear.md", "The grouped execution path preserves derivatives with respect to external weights, including repeated input and output irrep blocks."),
("docs/guide/performance.md", "The repository's `evals/` harness does this consistently across Torch CPU, general MLX, and kernel-enabled MLX."),
("docs/guide/performance.md", "Disabling a kernel changes execution, not the mathematical operation."),
("docs/guide/performance.md", "Dense shared-weight tensor products deliberately keep their parameter vector compact."),
("docs/guide/performance.md", "The shipped `gate_points_2102` model currently has no model-level kernel toggle"),
("evals/README.md", "Every result row also records the selected execution path."),
]
for path, excerpt in target_refs:
text = Path(path).read_text()
print(f"{path}:{path}:{excerpt!r}: FOUND" if excerpt in text else f"{path}:{path}:{excerpt!r}: MISSING")
PYRepository: lamalab-org/e3nn_mlx
Length of output: 2517
Fix the verbatim documentation contract excerpts.
DOCUMENTED_SOURCE_CONTRACTS currently contains literal substrings that do not exist in the referenced docs, so the test will fail under the exact-substring contract. For example, docs/guide/tensor_products.md only contains Shared weights remain one-dimensional during execution. before mentioning MLX broadcasting; several docs/guide/linear.md, docs/guide/performance.md, and evals/README.md quoted excerpts also need to match the actual wording/punctuation.
🤖 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 282 - 341, Update the
DOCUMENTED_SOURCE_CONTRACTS entries in test_documentation_contract.py so every
quoted excerpt exactly matches the current wording and punctuation in its
referenced documentation file. Inspect the corresponding sections of
tensor_products.md, linear.md, performance.md, and evals/README.md, then replace
only the stale literal substrings while preserving each associated test
reference and contract intent.
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
Refine performance evaluation and numerical parity infrastructure while tightening tensor product and Linear semantics and Metal kernel dispatch, adding three-way Torch CPU/MLX benchmark reporting and randomized cross-framework parity suites.
New Features:
Bug Fixes:
Enhancements:
Tests:
Summary by CodeRabbit
Bug Fixes
Linearweight-path ordering and grouped execution to preserve gradients and correct external-weight interpretation.TensorProductvalidation, dtype consistency, shared/per-sample broadcasting, and instruction accumulation for repeated/unweighted cases.TensorSquareuse_custom_kernelpropagation.Documentation
Linearsemantics/compatibility and clarified migration, kernel limitations, and benchmark/evaluation workflow.Testing