Types.

Author: ameya98

tests/test_extract_irreps.py

@@ -8,7 +8,7 @@
 
 
 @pytest.mark.parametrize("irreps_in", ["0e + 1o + 2e", "2x0e + 1o + 2x2e + 1x1o"])
-def test_equivariance(irreps_in):
+def test_equivariance(irreps_in: str):
     irreps_in = e3nn.o3.Irreps(irreps_in)
     layer = ExtractIrreps(irreps_in, "0e")
     e3nn.util.test.assert_equivariant(

tests/test_layer_norm.py

@@ -9,7 +9,7 @@
 
 
 @pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
-def test_equivariance(irreps_in):
+def test_equivariance(irreps_in: str):
     irreps_in = e3nn.o3.Irreps(irreps_in)
     layer = LayerNorm(irreps_in)
     e3nn.util.test.assert_equivariant(

tests/test_scaling.py

@@ -8,7 +8,7 @@
 
 
 @pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
-def test_equivariance(irreps_in):
+def test_equivariance(irreps_in: str):
     irreps_in = e3nn.o3.Irreps(irreps_in)
     layer = ScaleIrreps(irreps_in)
     irreps_weight = e3nn.o3.Irreps(f"{layer.irreps_in.num_irreps}x0e")
@@ -20,7 +20,7 @@ def test_equivariance(irreps_in):
 
 
 @pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
-def test_scale_irreps_by_one(irreps_in):
+def test_scale_irreps_by_one(irreps_in: str):
     irreps_in = e3nn.o3.Irreps(irreps_in)
     layer = ScaleIrreps(irreps_in)
     assert layer.irreps_in == irreps_in