Address benchmark and parity review findings

Author: r-fedorov

e3nn_mlx/ops_tp.py

@@ -1307,6 +1307,7 @@ def __init__(
         internal_weights: bool | None = None,
         shared_weights: bool = True,
         compile_left_right: bool = True,
+        use_custom_kernel: bool = True,
     ) -> None:
         irreps_in = Irreps(irreps_in).simplify()
         parsed_filter = None if filter_ir_out is None else [Irrep.parse(ir) for ir in filter_ir_out]
@@ -1324,6 +1325,7 @@ def __init__(
                 internal_weights=False,
                 shared_weights=True,
                 compile_left_right=compile_left_right,
+                use_custom_kernel=use_custom_kernel,
             )
             self._execution_irreps_out = self.irreps_out
             grouped_out, index_groups = _group_output_irreps(self._execution_irreps_out)
@@ -1351,6 +1353,7 @@ def __init__(
                 internal_weights=internal_weights,
                 shared_weights=shared_weights,
                 compile_left_right=compile_left_right,
+                use_custom_kernel=use_custom_kernel,
             )
             self._execution_irreps_out = self.irreps_out
         self.irreps_in = irreps_in

evals/mlx_cases.py

@@ -26,6 +26,12 @@ def _imports():
     return mx, nn, o3, scatter_sum
 
 
+def _metal_available() -> bool:
+    from e3nn_mlx.compat import mlx_metal_available
+
+    return mlx_metal_available()
+
+
 def backend_metadata() -> dict[str, Any]:
     mx, _, _, _ = _imports()
     return {
@@ -125,7 +131,15 @@ def forward(value):
         compile_train=_compiled(value_grad, vectors),
         dispatch=(
             "metal-spherical-harmonics"
-            if _USE_CUSTOM_KERNELS and config["lmax"] <= 4
+            if (
+                _USE_CUSTOM_KERNELS
+                and _metal_available()
+                and vectors.ndim == 2
+                and vectors.shape[0] > 0
+                and vectors.dtype == mx.float32
+                and degrees
+                and max(degrees) <= 4
+            )
             else (
                 "general-mlx (kernel fallback)"
                 if _USE_CUSTOM_KERNELS
@@ -301,7 +315,19 @@ def forward(source):
         train=lambda: value_grad(values),
         compile_train=_compiled(value_grad, values),
         dispatch=(
-            "metal-scatter-sum" if _USE_CUSTOM_KERNELS else "general-mlx"
+            "metal-scatter-sum"
+            if (
+                _USE_CUSTOM_KERNELS
+                and _metal_available()
+                and values.ndim >= 2
+                and values.shape[0] > 0
+                and values.dtype == mx.float32
+            )
+            else (
+                "general-mlx (kernel fallback)"
+                if _USE_CUSTOM_KERNELS
+                else "general-mlx"
+            )
         ),
     )
 
@@ -332,9 +358,7 @@ def _model_dispatch(*, supports_kernels: bool) -> str:
         return "general-mlx"
     if not supports_kernels:
         return "general-mlx (model has no kernel toggle)"
-    from e3nn_mlx.compat import mlx_metal_available
-
-    if not mlx_metal_available():
+    if not _metal_available():
         return "general-mlx (kernel fallback)"
     return "mixed-model-kernels"
 

evals/randomized_core_parity.py

@@ -722,11 +722,6 @@ def _torch_worker(cases: list[dict[str, Any]]) -> dict[str, Any]:
 
 
 def _mlx_variants(case: dict[str, Any]) -> tuple[tuple[str, bool, bool], ...]:
-    if (
-        case["family"] == "tensor_product_wrappers"
-        and case["kind"] == "tensor_square"
-    ):
-        return (("default", True, True),)
     if case["family"] in {"spherical_harmonics", "scatter", "tensor_product_wrappers"}:
         return (
             ("generic-eager", False, False),
@@ -891,6 +886,7 @@ def _mlx_module_output(
             module = o3.TensorSquare(
                 case["irreps_in1"],
                 compile_left_right=compiled,
+                use_custom_kernel=custom,
             )
             output = module(o3.IrrepsArray(module.irreps_in, left_values)).array
         metadata["custom_eligible"] = bool(

evals/torch_cases.py

@@ -81,6 +81,24 @@ def train():
     return train
 
 
+def _forward_with_fixed_radius_graph(
+    model_module,
+    module,
+    data,
+    edge_index,
+):
+    original_radius_graph = model_module.radius_graph
+
+    def fixed_radius_graph(*_args, **_kwargs):
+        return edge_index
+
+    model_module.radius_graph = fixed_radius_graph
+    try:
+        return module(data)
+    finally:
+        model_module.radius_graph = original_radius_graph
+
+
 def build_spherical_harmonics(config: dict[str, Any]) -> Task:
     torch, o3 = _imports()
     vectors = torch.randn(config["items"], 3, requires_grad=True)
@@ -141,7 +159,9 @@ def build_fully_connected_tensor_product(config: dict[str, Any]) -> Task:
     module = o3.FullyConnectedTensorProduct(irreps, irreps, irreps)
     left = torch.randn(config["items"], irreps.dim)
     right = torch.randn(config["items"], irreps.dim)
-    forward = lambda: module(left, right)
+    def forward():
+        return module(left, right)
+
     return _task(
         name="fully_connected_tensor_product",
         config=config,
@@ -212,7 +232,9 @@ def build_linear(config: dict[str, Any]) -> Task:
     irreps = o3.Irreps(spherical_irreps(config["mul"], config["lmax"]))
     module = o3.Linear(irreps, irreps)
     values = torch.randn(config["items"], irreps.dim)
-    forward = lambda: module(values)
+    def forward():
+        return module(values)
+
     return _task(
         name="linear",
         config=config,
@@ -267,7 +289,7 @@ def _model_graph(config, *, input_dim, node_attr_dim=0, edge_attr_dim=0):
 
 
 def build_gate_points_2102(config: dict[str, Any]) -> Task:
-    torch, o3 = _imports()
+    _, o3 = _imports()
     import e3nn.nn.models.gate_points_2102 as model_module
 
     irreps_in = o3.Irreps("4x0e")
@@ -297,14 +319,20 @@ def build_gate_points_2102(config: dict[str, Any]) -> Task:
         input_dim=irreps_in.dim,
         node_attr_dim=irreps_node_attr.dim,
     )
-    model_module.radius_graph = lambda _pos, _radius, _batch: edge_index
     data = {
         "pos": positions,
         "x": node_input,
         "z": node_attr,
         "batch": batch,
     }
-    forward = lambda: module(data)
+    def forward():
+        return _forward_with_fixed_radius_graph(
+            model_module,
+            module,
+            data,
+            edge_index,
+        )
+
     return _task(
         name="gate_points_2102",
         config=config,
@@ -315,7 +343,7 @@ def build_gate_points_2102(config: dict[str, Any]) -> Task:
 
 
 def build_v2106_simple_network(config: dict[str, Any]) -> Task:
-    torch, o3 = _imports()
+    _, o3 = _imports()
     import e3nn.nn.models.v2106.gate_points_networks as model_module
 
     irreps_in = o3.Irreps("4x0e")
@@ -334,9 +362,16 @@ def build_v2106_simple_network(config: dict[str, Any]) -> Task:
         config,
         input_dim=irreps_in.dim,
     )
-    model_module.radius_graph = lambda _pos, _radius, _batch: edge_index
     data = {"pos": positions, "x": node_input, "batch": batch}
-    forward = lambda: module(data)
+
+    def forward():
+        return _forward_with_fixed_radius_graph(
+            model_module,
+            module,
+            data,
+            edge_index,
+        )
+
     return _task(
         name="v2106_simple_network",
         config=config,
@@ -347,7 +382,7 @@ def build_v2106_simple_network(config: dict[str, Any]) -> Task:
 
 
 def build_v2106_attributed_network(config: dict[str, Any]) -> Task:
-    torch, o3 = _imports()
+    _, o3 = _imports()
     from e3nn.nn.models.v2106.gate_points_networks import (
         NetworkForAGraphWithAttributes,
     )
@@ -382,7 +417,9 @@ def build_v2106_attributed_network(config: dict[str, Any]) -> Task:
         "edge_index": edge_index,
         "batch": batch,
     }
-    forward = lambda: module(data)
+    def forward():
+        return module(data)
+
     return _task(
         name="v2106_attributed_network",
         config=config,

tests/test_evals.py

@@ -200,6 +200,12 @@ def test_mlx_benchmark_reports_actual_tensor_product_dispatch() -> None:
     dense = mlx_cases.build_fully_connected_tensor_product(
         {"items": 64, "mul": 8, "lmax": 2}
     )
+    harmonics = mlx_cases.build_spherical_harmonics(
+        {"items": 16, "lmax": 2}
+    )
+    scatter = mlx_cases.build_scatter_sum(
+        {"items": 16, "nodes": 4, "width": 3}
+    )
     mlx_cases.configure(use_custom_kernels=False)
     general = mlx_cases.build_fully_connected_tensor_product(
         {"items": 16, "mul": 8, "lmax": 2}
@@ -211,6 +217,16 @@ def test_mlx_benchmark_reports_actual_tensor_product_dispatch() -> None:
         else "general-mlx (kernel fallback)"
     )
     assert small.dispatch == expected
+    assert harmonics.dispatch == (
+        "metal-spherical-harmonics"
+        if mlx_metal_available()
+        else "general-mlx (kernel fallback)"
+    )
+    assert scatter.dispatch == (
+        "metal-scatter-sum"
+        if mlx_metal_available()
+        else "general-mlx (kernel fallback)"
+    )
     assert dense.dispatch == "general-mlx (kernel fallback)"
     assert general.dispatch == "general-mlx"
 
@@ -227,8 +243,16 @@ def test_mlx_benchmark_reports_non_metal_kernel_fallback(monkeypatch) -> None:
     task = mlx_cases.build_fully_connected_tensor_product(
         {"items": 16, "mul": 8, "lmax": 2}
     )
+    harmonics = mlx_cases.build_spherical_harmonics(
+        {"items": 16, "lmax": 2}
+    )
+    scatter = mlx_cases.build_scatter_sum(
+        {"items": 16, "nodes": 4, "width": 3}
+    )
     model_dispatch = mlx_cases._model_dispatch(supports_kernels=True)
     mlx_cases.configure(use_custom_kernels=False)
 
     assert task.dispatch == "general-mlx (kernel fallback)"
+    assert harmonics.dispatch == "general-mlx (kernel fallback)"
+    assert scatter.dispatch == "general-mlx (kernel fallback)"
     assert model_dispatch == "general-mlx (kernel fallback)"

tests/test_metal_kernels.py

@@ -211,7 +211,7 @@ def apply(module, first, second, value):
         mx.ones_like(right) * -0.03,
         mx.ones_like(weight) * 0.11,
     )
-    with pytest.raises(ValueError, match="Not implemented for CustomKernel"):
+    with pytest.raises(ValueError, match="CustomKernel"):
         mx.jvp(
             lambda first, second, value: apply(
                 kernel, first, second, value

tests/test_randomized_tensor_product_parity_harness.py

@@ -74,18 +74,20 @@ def test_comparison_reports_every_execution_variant_and_retains_failures():
         ],
     }
     variants = []
-    for index, (name, _, _) in enumerate(MLX_VARIANTS):
+    for name, _, 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
+                ),
             }
         )
     mlx_result = {

tests/test_tensor_product_compatibility_qualification.py

@@ -286,22 +286,31 @@ def test_shared_tensor_product_weights_are_not_materialized_per_item(
         use_custom_kernel=False,
     )
     left = _array(product.irreps_in1, mx.ones((8, product.irreps_in1.dim)))
-    right = _array(product.irreps_in2, mx.ones((8, product.irreps_in2.dim)))
+    right = _array(product.irreps_in2, mx.ones((1, product.irreps_in2.dim)))
     shared_weight = product.weight
-    materialized_shapes = []
+    broadcasts = []
     original_broadcast_to = mx.broadcast_to
 
     def record_broadcast(array, shape, *args, **kwargs):
-        if array is shared_weight:
-            materialized_shapes.append(tuple(shape))
+        broadcasts.append((tuple(array.shape), 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 == []
+    assert broadcasts
+    weight_broadcasts = [
+        target_shape
+        for source_shape, target_shape in broadcasts
+        if source_shape == tuple(shared_weight.shape)
+        or (
+            source_shape
+            and source_shape[-1] == product.weight_numel
+        )
+    ]
+    assert weight_broadcasts == []
 
 
 @pytest.mark.mlx

tests/test_tensor_product_wrappers.py

@@ -49,6 +49,15 @@ def test_tensor_square_symbolic_full_mode() -> None:
     assert {inst.mode for inst in tp.instructions} == {"uvu<v", "uuu"}
 
 
+def test_tensor_square_propagates_custom_kernel_setting() -> None:
+    general = TensorSquare("2x0e", use_custom_kernel=False)
+    requested = TensorSquare("2x0e", use_custom_kernel=True)
+
+    assert general.use_custom_kernel is False
+    assert general._metal_operation is None
+    assert requested.use_custom_kernel is True
+
+
 def test_tensor_square_symbolic_fully_connected_mode() -> None:
     tp = TensorSquare("2x0e", "1x0e", internal_weights=False)
     assert str(tp.irreps_out) == "0e"

tests/test_upstream_o3_linear_norm.py

@@ -158,6 +158,7 @@ def test_grouped_linear_external_weight_vjp_matches_blockwise_formula() -> None:
     module = o3.Linear(
         "3x2e + 3x3e + 2x2e",
         "3x2e + 3x3e + 3x2e",
+        instructions=[(0, 0), (1, 1), (2, 2)],
         internal_weights=False,
         path_normalization="path",
     )