Port message-passing convention fix to warp TensorNet (#766)

The warp embedding (_embedding_warp.py) was aggregating radial messages
onto destination nodes via col_data/col_indptr, the old convention
corrected on main in PRs #758/#759 for the PyG/DGL implementations.
Switch to row_data/row_indptr so warp matches the corrected aggregation
onto source nodes, and add a parity test that loads the pretrained
TensorNet-PES-MatPES-PBE-2025.2 weights into both warp and non-warp
TensorNet instances and asserts identical outputs. Re-calibrate the
existing warp regression-check expected values affected by the fix.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: shyuep

src/matgl/layers/_embedding_warp.py

@@ -115,8 +115,8 @@ def forward(
         edge_weight: torch.Tensor,
         edge_vec: torch.Tensor,
         edge_attr: torch.Tensor,
-        col_data: torch.Tensor,
-        col_indptr: torch.Tensor,
+        row_data: torch.Tensor,
+        row_indptr: torch.Tensor,
     ) -> torch.Tensor:
         """Forward pass.
 
@@ -126,8 +126,8 @@ def forward(
             edge_weight: Edge weights (distances), shape (num_edges,)
             edge_vec: Edge vectors, shape (num_edges, 3)
             edge_attr: Edge attributes (RBF), shape (num_edges, num_rbf)
-            col_data: CSC col data for destination aggregation, shape (num_edges,)
-            col_indptr: CSC col indptr for destination aggregation, shape (num_nodes+1,)
+            row_data: CSR row data for source aggregation, shape (num_edges,)
+            row_indptr: CSR row indptr for source aggregation, shape (num_nodes+1,)
 
         Returns:
             X: Tensor representation, shape (num_nodes, 3, 3, units)
@@ -143,7 +143,7 @@ def forward(
         edge_attr_processed = edge_attr.view(-1, 3, self.units) * C.view(-1, 1, 1) * Zij.view(-1, 1, self.units)
 
         edge_vec_norm = edge_vec / torch.norm(edge_vec, dim=1, keepdim=True).clamp(min=1e-6)
-        I, A, S = fn_radial_message_passing(edge_vec_norm, edge_attr_processed, col_data, col_indptr)  # noqa: E741
+        I, A, S = fn_radial_message_passing(edge_vec_norm, edge_attr_processed, row_data, row_indptr)  # noqa: E741
 
         X = fn_compose_tensor(I, A, S)  # (num_nodes, 3, 3, units)
 

src/matgl/models/_tensornet_pyg.py

@@ -283,7 +283,7 @@ def forward_features(
                 col_indices,
                 col_indptr,
             ) = graph_transform(edge_index.int(), z.shape[0])  # type: ignore[union-attr]
-            X = self.tensor_embedding(z, edge_index, bond_dist, bond_vec, edge_attr, col_data, col_indptr)
+            X = self.tensor_embedding(z, edge_index, bond_dist, bond_vec, edge_attr, row_data, row_indptr)
             fea_dict["embedding"] = X
             for i, layer in enumerate(self.layers):
                 X = layer(

tests/models/test_tensornet_warp.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -12,6 +13,7 @@
     pytest.skip("Skipping PYG tests", allow_module_level=True)
 
 from matgl.models._tensornet_pyg import TensorNet, _warp_available
+from matgl.utils.io import _get_file_paths
 
 if not _warp_available:
     pytest.skip("Skipping warp tests: nvalchemiops not installed", allow_module_level=True)
@@ -23,11 +25,11 @@ def test_model(graph_MoS_pyg):
 
     # Optional regression-check values
     EXPECTED = {
-        "swish": torch.tensor(0.0813),
-        "tanh": torch.tensor(-0.0189),
-        "sigmoid": torch.tensor(0.0353),
-        "softplus2": torch.tensor(0.1164),
-        "softexp": torch.tensor(0.1148),
+        "swish": torch.tensor(0.0827),
+        "tanh": torch.tensor(-0.0258),
+        "sigmoid": torch.tensor(0.0360),
+        "softplus2": torch.tensor(0.1165),
+        "softexp": torch.tensor(0.1100),
     }
 
     _, graph, _ = graph_MoS_pyg
@@ -78,7 +80,7 @@ def test_model_intensive(graph_MoS_pyg):
     graph.pos = graph.frac_coords @ lat[0]
     model = TensorNet(element_types=["Mo", "S"], is_intensive=True)
     output = model(g=graph)
-    assert torch.allclose(output, torch.tensor([-0.0897]), atol=1e-4)
+    assert torch.allclose(output, torch.tensor([-0.0906]), atol=1e-4)
 
 
 def test_model_intensive_with_weighted_atom(graph_MoS_pyg):
@@ -88,7 +90,7 @@ def test_model_intensive_with_weighted_atom(graph_MoS_pyg):
     graph.pos = graph.frac_coords @ lat[0]
     model = TensorNet(element_types=["Mo", "S"], is_intensive=True, readout_type="weighted_atom")
     output = model(g=graph)
-    assert torch.allclose(output, torch.tensor([-0.0217]), atol=1e-4)
+    assert torch.allclose(output, torch.tensor([-0.0210]), atol=1e-4)
 
 
 def test_model_intensive_with_ReduceReadOut(graph_MoS_pyg):
@@ -98,7 +100,7 @@ def test_model_intensive_with_ReduceReadOut(graph_MoS_pyg):
     graph.pos = graph.frac_coords @ lat[0]
     model = TensorNet(is_intensive=True, readout_type="reduce_atom")
     output = model(g=graph)
-    assert torch.allclose(output, torch.tensor([-0.1045]), atol=1e-4)
+    assert torch.allclose(output, torch.tensor([-0.1075]), atol=1e-4)
 
 
 def test_model_intensive_with_classification(graph_MoS_pyg):
@@ -122,9 +124,9 @@ def test_backward(graph_MoS_pyg):
 
     EXPECTED_CELL_GRAD = torch.tensor(
         [
-            [-0.000967, 0.000000, 0.000000],
-            [0.000000, -0.000967, 0.000000],
-            [0.000000, 0.000000, -0.000967],
+            [-0.000909, 0.000000, 0.000000],
+            [0.000000, -0.000909, 0.000000],
+            [0.000000, 0.000000, -0.000909],
         ]
     )
 
@@ -150,9 +152,9 @@ def test_double_backward(graph_MoS_pyg):
 
     EXPECTED_CELL_GRAD2 = torch.tensor(
         [
-            [-0.000010, -0.000000, -0.000000],
-            [-0.000000, -0.000010, -0.000000],
-            [-0.000000, -0.000000, -0.000010],
+            [-0.0000037, -0.000000, -0.000000],
+            [-0.000000, -0.0000037, -0.000000],
+            [-0.000000, -0.000000, -0.0000037],
         ]
     )
 
@@ -172,3 +174,42 @@ def test_double_backward(graph_MoS_pyg):
     loss.backward()
 
     assert torch.allclose(cell.grad, EXPECTED_CELL_GRAD2, atol=1e-6)
+
+
+def _build_pair_from_pretrained(repo_id: str) -> tuple[TensorNet, TensorNet]:
+    """Build a (warp, non-warp) pair of TensorNet models loaded with identical pretrained weights."""
+    fpaths = _get_file_paths(Path(repo_id))
+    map_location = "cpu" if not torch.cuda.is_available() else None
+    state = torch.load(fpaths["state.pt"], map_location=map_location, weights_only=False)
+    init_blob = torch.load(fpaths["model.pt"], map_location=map_location, weights_only=False)
+    inner_init_args = dict(init_blob["model"]["init_args"])
+
+    inner_state = {k[len("model.") :]: v for k, v in state.items() if k.startswith("model.")}
+
+    model_warp = TensorNet(**{**inner_init_args, "use_warp": True})
+    model_pyg = TensorNet(**{**inner_init_args, "use_warp": False})
+    model_warp.load_state_dict(inner_state, strict=False)
+    model_pyg.load_state_dict(inner_state, strict=False)
+    model_warp.eval()
+    model_pyg.eval()
+    return model_warp, model_pyg
+
+
+def test_warp_pyg_parity_pretrained(MoS):
+    """Warp and non-warp TensorNet must produce identical outputs from the same pretrained weights."""
+    model_warp, model_pyg = _build_pair_from_pretrained("materialyze/TensorNet-PES-MatPES-PBE-2025.2")
+
+    from matgl.ext._pymatgen_pyg import Structure2Graph
+
+    converter = Structure2Graph(element_types=model_pyg.element_types, cutoff=model_pyg.cutoff)
+    g, lat, _ = converter.get_graph(MoS)
+    g.pbc_offshift = torch.matmul(g.pbc_offset, lat[0])
+    g.pos = g.frac_coords @ lat[0]
+
+    with torch.no_grad():
+        out_warp = model_warp(g=g)
+        out_pyg = model_pyg(g=g)
+
+    assert torch.allclose(out_warp, out_pyg, atol=1e-5, rtol=1e-5), (
+        f"warp={out_warp.detach().cpu()} vs pyg={out_pyg.detach().cpu()}"
+    )