Fix cross-platform CI dispatch checks

Author: r-fedorov

evals/mlx_cases.py

@@ -332,6 +332,10 @@ 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():
+        return "general-mlx (kernel fallback)"
     return "mixed-model-kernels"
 
 

tests/test_evals.py

@@ -163,6 +163,7 @@ def test_backend_builders_match_the_documented_case_set() -> None:
 @pytest.mark.mlx
 def test_model_benchmarks_report_kernel_scope() -> None:
     from evals import mlx_cases
+    from e3nn_mlx.compat import mlx_metal_available
 
     config = {
         "nodes": 8,
@@ -178,13 +179,19 @@ def test_model_benchmarks_report_kernel_scope() -> None:
     mlx_cases.configure(use_custom_kernels=False)
 
     assert gate.dispatch == "general-mlx (model has no kernel toggle)"
-    assert simple.dispatch == "mixed-model-kernels"
-    assert attributed.dispatch == "mixed-model-kernels"
+    expected = (
+        "mixed-model-kernels"
+        if mlx_metal_available()
+        else "general-mlx (kernel fallback)"
+    )
+    assert simple.dispatch == expected
+    assert attributed.dispatch == expected
 
 
 @pytest.mark.mlx
 def test_mlx_benchmark_reports_actual_tensor_product_dispatch() -> None:
     from evals import mlx_cases
+    from e3nn_mlx.compat import mlx_metal_available
 
     mlx_cases.configure(use_custom_kernels=True)
     small = mlx_cases.build_fully_connected_tensor_product(
@@ -198,6 +205,30 @@ def test_mlx_benchmark_reports_actual_tensor_product_dispatch() -> None:
         {"items": 16, "mul": 8, "lmax": 2}
     )
 
-    assert small.dispatch == "metal-scalar-paths"
+    expected = (
+        "metal-scalar-paths"
+        if mlx_metal_available()
+        else "general-mlx (kernel fallback)"
+    )
+    assert small.dispatch == expected
     assert dense.dispatch == "general-mlx (kernel fallback)"
     assert general.dispatch == "general-mlx"
+
+
+@pytest.mark.mlx
+def test_mlx_benchmark_reports_non_metal_kernel_fallback(monkeypatch) -> None:
+    import e3nn_mlx.compat as compat
+    import e3nn_mlx.ops_tp as tp_module
+    from evals import mlx_cases
+
+    monkeypatch.setattr(compat, "mlx_metal_available", lambda: False)
+    monkeypatch.setattr(tp_module, "mlx_metal_available", lambda: False)
+    mlx_cases.configure(use_custom_kernels=True)
+    task = mlx_cases.build_fully_connected_tensor_product(
+        {"items": 16, "mul": 8, "lmax": 2}
+    )
+    model_dispatch = mlx_cases._model_dispatch(supports_kernels=True)
+    mlx_cases.configure(use_custom_kernels=False)
+
+    assert task.dispatch == "general-mlx (kernel fallback)"
+    assert model_dispatch == "general-mlx (kernel fallback)"

tests/test_v2106_convolution.py

@@ -154,11 +154,17 @@ def test_v2106_convolution_compilation_and_cached_reuse() -> None:
     expected = module.forward_arrays(*values)
     mx.eval(expected)
     compiled = mx.compile(module.forward_arrays)
-    assert _max_abs(compiled(*values) - expected) < 5e-5
+    actual = compiled(*values)
+    mx.eval(actual)
+    assert bool(mx.allclose(actual, expected, atol=5e-5, rtol=2e-6))
     changed = (values[0] * 0.9, values[1], *values[2:])
     changed_expected = module.forward_arrays(*changed)
     mx.eval(changed_expected)
-    assert _max_abs(compiled(*changed) - changed_expected) < 5e-5
+    changed_actual = compiled(*changed)
+    mx.eval(changed_actual)
+    assert bool(
+        mx.allclose(changed_actual, changed_expected, atol=5e-5, rtol=2e-6)
+    )
 
 
 @pytest.mark.mlx