assert torch.allclose -> torch.testing.assert_close

Author: kleinhenz

tests/test_attention.py

@@ -87,7 +87,7 @@ def test_compile(model_name, causal, request, irreps_in, irreps_sh, edge_attr_di
     out = model(node_attr, edge_index, edge_attr, edge_sh, mask=mask)
     compiled_out = compiled_model(node_attr, edge_index, edge_attr, edge_sh, mask=mask)
 
-    assert torch.allclose(out, compiled_out)
+    torch.testing.assert_close(out, compiled_out)
 
 
 @pytest.mark.parametrize("model_name", ["singlehead_attention", "multihead_attention"])
@@ -124,20 +124,22 @@ def test_causal_vs_non_causal_attention(
     causal_mask = edge_index[0] <= edge_index[1]
     non_causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=None)
     causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=causal_mask)
-    assert torch.allclose(non_causal_out, causal_out)
+    torch.testing.assert_close(non_causal_out, causal_out)
 
     # Check that the outputs are the same for the nodes that do not have any causal edges.
     edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]])
     causal_mask = edge_index[0] <= edge_index[1]
     non_causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=None)
     causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=causal_mask)
-    assert not torch.allclose(non_causal_out[:1], causal_out[:1])
-    assert torch.allclose(non_causal_out[1:], causal_out[1:])
+    with pytest.raises(AssertionError):
+        torch.testing.assert_close(non_causal_out[:1], causal_out[:1])
+    torch.testing.assert_close(non_causal_out[1:], causal_out[1:])
 
     # Check that the outputs are the same for the nodes that do not have any causal edges.
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 4]])
     causal_mask = edge_index[0] <= edge_index[1]
     non_causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=None)
     causal_out = model(node_attr, edge_index, edge_attr, edge_sh, mask=causal_mask)
-    assert not torch.allclose(non_causal_out[:2], causal_out[:2])
-    assert torch.allclose(non_causal_out[2:], causal_out[2:])
+    with pytest.raises(AssertionError):
+        torch.testing.assert_close(non_causal_out[:2], causal_out[:2])
+    torch.testing.assert_close(non_causal_out[2:], causal_out[2:])

tests/test_equivariance.py

@@ -72,7 +72,7 @@ def test_conv(conv):
         layer = conv(irreps_in, irreps_in, irreps_sh, edge_attr_dim=edge_attr_dim)
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 @pytest.mark.parametrize("conv", CONV_LAYERS)
@@ -89,7 +89,7 @@ def test_gated_conv(conv):
         layer = Gated(wrapped, irreps_in=irreps_in, irreps_out=irreps_in)
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 @pytest.mark.parametrize("conv", CONV_LAYERS)
@@ -108,7 +108,7 @@ def test_conv_block(conv):
         )
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 @pytest.mark.parametrize("conv", CONV_LAYERS)
@@ -132,7 +132,7 @@ def test_attention(conv):
         )
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 @pytest.mark.parametrize("conv", [Conv, SeparableConv])
@@ -158,7 +158,7 @@ def test_multihead_attention(conv):
         )
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 def test_layer_norm():
@@ -174,7 +174,7 @@ def test_layer_norm():
         out_1 = layer(x @ D.T)
         out_2 = layer(x) @ D.T
 
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 def test_equivariant_mlp():
@@ -194,7 +194,7 @@ def test_equivariant_mlp():
         out_1 = layer(x @ D.T)
         out_2 = layer(x) @ D.T
 
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)
 
 
 def test_transformer():
@@ -214,4 +214,4 @@ def test_transformer():
         )
 
         out_1, out_2 = apply_layer_rotation(layer)
-        assert torch.allclose(out_1, out_2, atol=1e-10)
+        torch.testing.assert_close(out_1, out_2, atol=1e-10, rtol=1e-10)

tests/test_extract_irreps.py

@@ -32,15 +32,15 @@ def test_extract_irreps_simple():
 
     layer = ExtractIrreps(irreps_in, "0e")
     output = layer(input)
-    assert torch.allclose(output, torch.as_tensor([1.0]))
+    torch.testing.assert_close(output, torch.as_tensor([1.0]))
 
     layer = ExtractIrreps(irreps_in, "1o")
     output = layer(input)
-    assert torch.allclose(output, torch.as_tensor([2.0, 3.0, 4.0]))
+    torch.testing.assert_close(output, torch.as_tensor([2.0, 3.0, 4.0]))
 
     layer = ExtractIrreps(irreps_in, "2e")
     output = layer(input)
-    assert torch.allclose(output, torch.as_tensor([5.0, 6.0, 7.0, 8.0, 9.0]))
+    torch.testing.assert_close(output, torch.as_tensor([5.0, 6.0, 7.0, 8.0, 9.0]))
 
 
 def test_extract_irreps_multiplicity():
@@ -50,8 +50,8 @@ def test_extract_irreps_multiplicity():
 
     layer = ExtractIrreps(irreps_in, "0e")
     output = layer(input)
-    assert torch.allclose(output, torch.as_tensor([1.0, 5.0, 6.0]))
+    torch.testing.assert_close(output, torch.as_tensor([1.0, 5.0, 6.0]))
 
     layer = ExtractIrreps(irreps_in, "1o")
     output = layer(input)
-    assert torch.allclose(output, torch.as_tensor([2.0, 3.0, 4.0, 7.0, 8.0, 9.0]))
+    torch.testing.assert_close(output, torch.as_tensor([2.0, 3.0, 4.0, 7.0, 8.0, 9.0]))

tests/test_fused_conv.py

@@ -82,4 +82,4 @@ def radial_nn(edge_attr_dim: int, num_elements: int) -> nn.Module:
     out = layer(node_attr, edge_index, edge_attr, edge_sh)
     out_fused = fused_layer(node_attr, edge_index, edge_attr, edge_sh)
 
-    assert torch.allclose(out, out_fused, rtol=1e-3)
+    torch.testing.assert_close(out, out_fused, rtol=1e-3, atol=1e-5)

tests/test_layer_norm.py

@@ -41,7 +41,7 @@ def test_layer_norm_compiled(irreps_in: str, seed: int, batch_size: int = 8):
     output = layer(input)
     output_compiled = layer_compiled(input)
 
-    assert torch.allclose(output, output_compiled)
+    torch.testing.assert_close(output, output_compiled)
 
 
 @pytest.mark.parametrize(
@@ -57,8 +57,13 @@ def test_layer_norm(irreps_in: str):
     output = layer(input)
 
     for mul, ir, field in unpack_irreps(output, irreps_in):
-        sq_norms = field.norm(dim=-1, keepdim=True).pow(2).sum(dim=-1).mean(dim=-1)
+        sq_norms = (
+            field.norm(dim=-1, keepdim=True)
+            .pow(2)
+            .sum(dim=-1)
+            .mean(dim=-1, keepdim=True)
+        )
         if ir.l == 0 and ir.p == 1 and mul == 1:
-            assert torch.allclose(sq_norms, torch.as_tensor([0.0]))
+            torch.testing.assert_close(sq_norms, torch.as_tensor([0.0]))
         else:
-            assert torch.allclose(sq_norms, torch.as_tensor([1.0]))
+            torch.testing.assert_close(sq_norms, torch.as_tensor([1.0]))

tests/test_pack_unpack.py

@@ -100,7 +100,7 @@ def test_inverse(irreps_in: str, factor: int, batch_size: int = 5):
     output = layer(input)
     recovered = inv_layer(output)
 
-    assert torch.allclose(input, recovered)
+    torch.testing.assert_close(input, recovered)
 
 
 @pytest.mark.parametrize(
@@ -116,7 +116,7 @@ def test_axis_to_mul_compiled(irreps_in: str, factor: int, batch_size: int = 5):
     layer = AxisToMul(irreps_in, factor)
     layer_compiled = torch.compile(layer, fullgraph=True)
 
-    assert torch.allclose(layer(input), layer_compiled(input))
+    torch.testing.assert_close(layer(input), layer_compiled(input))
 
 
 @pytest.mark.parametrize(
@@ -132,4 +132,4 @@ def test_mul_to_axis_compiled(irreps_in: str, factor: int, batch_size: int = 5):
     layer = MulToAxis(irreps_in, factor)
     layer_compiled = torch.compile(layer, fullgraph=True)
 
-    assert torch.allclose(layer(input), layer_compiled(input))
+    torch.testing.assert_close(layer(input), layer_compiled(input))

tests/test_scaling.py

@@ -30,7 +30,7 @@ def test_scale_irreps_by_one(irreps_in: str):
     weight = torch.ones(irreps_in.num_irreps)
     output = layer(input, weight)
 
-    assert torch.allclose(input, output)
+    torch.testing.assert_close(input, output)
 
 
 @pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
@@ -46,4 +46,4 @@ def test_scale_irreps_random(irreps_in: str):
 
     norm = e3nn.o3.Norm(irreps_in)
     factor = norm(output) / norm(input)
-    assert torch.allclose(factor, torch.abs(weight))
+    torch.testing.assert_close(factor, torch.abs(weight))