Validate periodic attention inputs early

Author: allaffa

hydragnn/globalAtt/equivariant_transformer.py

@@ -186,6 +186,12 @@ def forward_periodic_attention(
         """
         if batch.ndim != 1 or batch.shape[0] != node_features.shape[0]:
             raise ValueError("batch must contain one graph identifier per node")
+        if batch.dtype != torch.long:
+            raise TypeError("batch must have dtype torch.long")
+        if node_features.device != positions.device or positions.device != batch.device:
+            raise ValueError("features, positions, and batch must share a device")
+        if node_features.dtype != positions.dtype:
+            raise ValueError("features and positions must have the same dtype")
         normalized = self.attention_norm(node_features)
         sources = build_periodic_attention_sources(
             normalized, positions, batch, cell, pbc, replication

tests/test_equivariant_periodic_attention.py

@@ -148,6 +148,28 @@ def test_periodic_attention_keeps_batched_graphs_independent():
     torch.testing.assert_close(together, separate)
 
 
+def test_periodic_attention_validates_inputs_before_building_sources():
+    layer = EquivariantTransformerLayer(
+        "4x0e", heads=1, require_tensor_coupling=False
+    ).double()
+    features, positions, batch, cell, pbc = _periodic_inputs()
+
+    with pytest.raises(TypeError, match="batch must have dtype torch.long"):
+        layer.forward_periodic_attention(
+            features, positions, batch.to(torch.int32), cell, pbc
+        )
+    with pytest.raises(
+        ValueError, match="features and positions must have the same dtype"
+    ):
+        layer.forward_periodic_attention(features, positions.float(), batch, cell, pbc)
+    with pytest.raises(
+        ValueError, match="features, positions, and batch must share a device"
+    ):
+        layer.forward_periodic_attention(
+            features, positions, batch.to("meta"), cell, pbc
+        )
+
+
 def test_periodic_configuration_validates_replication_convention():
     config = {
         "global_attn_engine": "EquivariantTransformer",