Stabilize compiled model comparisons

Author: r-fedorov

tests/test_gate_points_2102_network.py

@@ -150,13 +150,18 @@ def forward(pos, node_input, node_attr, graph_batch, edge_src, edge_dst):
         ).array
 
     expected = forward(data["pos"], data["x"], data["z"], batch, edges[0], edges[1])
+    mx.eval(expected)
     compiled = mx.compile(forward)
     actual = compiled(data["pos"], data["x"], data["z"], batch, edges[0], edges[1])
     assert _max_abs(actual - expected) < 8e-5
     changed_x = data["x"] * 0.8
+    changed_expected = forward(
+        data["pos"], changed_x, data["z"], batch, edges[0], edges[1]
+    )
+    mx.eval(changed_expected)
     assert _max_abs(
         compiled(data["pos"], changed_x, data["z"], batch, edges[0], edges[1])
-        - forward(data["pos"], changed_x, data["z"], batch, edges[0], edges[1])
+        - changed_expected
     ) < 8e-5
 
 

tests/test_v2106_convolution.py

@@ -105,10 +105,14 @@ def test_v2106_convolution_compilation_and_cached_reuse() -> None:
     module = _module()
     _activate_alpha(module)
     values = _inputs(module)
+    expected = module.forward_arrays(*values)
+    mx.eval(expected)
     compiled = mx.compile(module.forward_arrays)
-    assert _max_abs(compiled(*values) - module.forward_arrays(*values)) < 5e-5
+    assert _max_abs(compiled(*values) - expected) < 5e-5
     changed = (values[0] * 0.9, values[1], *values[2:])
-    assert _max_abs(compiled(*changed) - module.forward_arrays(*changed)) < 5e-5
+    changed_expected = module.forward_arrays(*changed)
+    mx.eval(changed_expected)
+    assert _max_abs(compiled(*changed) - changed_expected) < 5e-5
 
 
 @pytest.mark.mlx

tests/test_v2106_message_passing.py

@@ -104,10 +104,14 @@ def test_v2106_message_passing_compilation_and_cached_reuse() -> None:
     module = _module()
     _activate_alpha(module)
     values = _inputs(module)
+    expected = module.forward_arrays(*values)
+    mx.eval(expected)
     compiled = mx.compile(module.forward_arrays)
-    assert _max_abs(compiled(*values) - module.forward_arrays(*values)) < 5e-4
+    assert _max_abs(compiled(*values) - expected) < 5e-4
     changed = (values[0] * 1.1, values[1], *values[2:])
-    assert _max_abs(compiled(*changed) - module.forward_arrays(*changed)) < 5e-4
+    changed_expected = module.forward_arrays(*changed)
+    mx.eval(changed_expected)
+    assert _max_abs(compiled(*changed) - changed_expected) < 5e-4
 
 
 @pytest.mark.mlx

tests/test_v2106_networks.py

@@ -243,11 +243,14 @@ def forward(pos, node_input, node_attr, edge_attr, graph_batch, edge_src, edge_d
             edges[1],
         )
     expected = forward(*args)
+    mx.eval(expected)
     compiled = mx.compile(forward)
     assert _max_abs(compiled(*args) - expected) < 2e-4
     changed = list(args)
     changed[1] = changed[1] * 0.9
-    assert _max_abs(compiled(*changed) - forward(*changed)) < 2e-4
+    changed_expected = forward(*changed)
+    mx.eval(changed_expected)
+    assert _max_abs(compiled(*changed) - changed_expected) < 2e-4
 
 
 @pytest.mark.mlx