propagate changes to tests and tutorial config

Author: cw-tan

configs/tutorial.yaml

@@ -2,7 +2,7 @@
 run: [train, test]
 
 cutoff_radius: 5.0
-model_type_names: [C, H, O, Cu]
+model_type_names: [C, H, O]
 chemical_species: ${model_type_names}
 
 data:
@@ -139,7 +139,7 @@ training_module:
 
     # === misc hyperparameters ===
     # average number of neighbors for edge sum normalization
-    avg_num_neighbors: ${training_data_stats:num_neighbors_mean}
+    avg_num_neighbors: ${training_data_stats:per_type_num_neighbors_mean}
 
     # per-type per-atom scales and shifts
     per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean}

tests/nn/test_contract_basic.py

@@ -57,31 +57,23 @@ def test_contract_jit(
             c_opt_mod.weights.copy_(c_base.weights)
         c_opt_mod = torch.jit.script(c_opt_mod)
 
-        def c_opt(x, y, idx, dim, w=None):
-            args = (x, y, idx, dim, w)
-            if w is None:
-                args = args[:-1]
-            return c_opt_mod(*args)
+        def c_opt(x, y, idx):
+            return c_opt_mod(x, y, idx)
 
-        batchdim = 17
-        scatter_dim = torch.tensor([batchdim], dtype=torch.long, device=device)
-        scatter_idxs = torch.arange(batchdim, device=device)
+        num_edges = 17
+        num_nodes = 11
+        scatter_idxs = torch.randint(0, num_nodes, (num_edges,), device=device)
         args_in = (
-            irreps_in1.randn(batchdim, mul, -1, device=device),
-            irreps_in2.randn(batchdim, mul, -1, device=device),
+            irreps_in1.randn(num_edges, mul, -1, device=device),
+            irreps_in2.randn(num_nodes, mul, -1, device=device),
             scatter_idxs,
-            scatter_dim,
-            torch.randn(
-                tuple(batchdim if e == -1 else e for e in c_base.weights.shape)
-            ),
         )
-        args_in = args_in[:-1]
 
         for c in (c_base, c_opt):
             assert_equivariant(
                 c,
                 args_in=args_in,
-                irreps_in=[irreps_in1, irreps_in2, None, None],
+                irreps_in=[irreps_in1, irreps_in2, None],
                 irreps_out=irreps_out,
                 # e3nn uses 1e-3, 1e-9
                 tolerance={torch.float32: 1e-3, torch.float64: 1e-8}[
@@ -164,7 +156,6 @@ def test_like_tp(
         # make input data
         batchdim = 1
         scatter_idxs = torch.arange(batchdim, device=device)
-        scatter_dim = torch.tensor([batchdim], dtype=torch.long, device=device)
         tensor1 = torch.randn(batchdim, mul, c.base_dim1, device=device)
         tensor2 = torch.randn(batchdim, mul, c.base_dim2, device=device)
 
@@ -193,9 +184,7 @@ def test_like_tp(
             weights_tp = weights_tp.T
         # else weights are just (u,)
         weights_tp = weights_tp.reshape(-1)
-        c_out = _strided_to_cat(
-            irreps_out, mul, c(tensor1, tensor2, scatter_idxs, scatter_dim)
-        )
+        c_out = _strided_to_cat(irreps_out, mul, c(tensor1, tensor2, scatter_idxs))
         tp_out = tp(
             _strided_to_cat(irreps_in1, mul, tensor1),
             _strided_to_cat(irreps_in2, mul, tensor2),

tests/nn/test_contract_kernels.py

@@ -97,16 +97,15 @@ def test_contract_kernel(
         )
         args_in = (
             irreps_in1.randn(num_edges, mul, -1, device=device),
-            irreps_in2.randn(num_edges, mul, -1, device=device),
+            irreps_in2.randn(num_atoms, mul, -1, device=device),
             scatter_idxs,
-            torch.tensor([num_atoms], dtype=torch.int64, device=device),
         )
 
         for c in (c_base, c_kernel):
             assert_equivariant(
                 c,
                 args_in=args_in,
-                irreps_in=[irreps_in1, irreps_in2, None, None],
+                irreps_in=[irreps_in1, irreps_in2, None],
                 irreps_out=irreps_out,
                 # e3nn uses 1e-3, 1e-9
                 tolerance={torch.float32: 1e-3, torch.float64: 1e-8}[