Extended Simformer and tests to handle non 3-dimensional data

Author: nMaax

sbi/neural_nets/net_builders/vector_field_nets.py

@@ -1269,6 +1269,12 @@ def forward(
             Vector field evaluation at the provided points
         """
 
+        if self.in_features == 1 and inputs.dim() == 2:
+            inputs = inputs.unsqueeze(-1)  # [B, T] -> [B, T, 1]
+            to_squeeze = True
+        else:
+            to_squeeze = False
+
         B, T, _ = inputs.shape
         device = inputs.device
 
@@ -1317,6 +1323,10 @@ def forward(
 
         # Output projection
         out = self.out_linear(h)  # [B, T, F]
+
+        if to_squeeze:
+            out = out.squeeze(-1)
+
         return out
 
 
@@ -1585,7 +1595,7 @@ def build_simformer_network(
     del batch_y  # Unused
     del embedding_net  # Unused
 
-    in_features = batch_x.shape[-1]
+    in_features = 1 if batch_x.dim() == 2 else batch_x.shape[-1]
     num_nodes = batch_x.shape[1]
 
     # Create the vector field network (Simformer)

tests/vector_field_nets_test.py

@@ -160,7 +160,7 @@ def test_vector_field_builders_shape_and_build(
 @pytest.mark.parametrize("time_emb_type", ["sinusoidal", "random_fourier"])
 @pytest.mark.parametrize("time_embedding_dim", [8, 16])
 @pytest.mark.parametrize("batch_dim", [1, 3])
-@pytest.mark.parametrize("input_dim", [(1, 1), (2, 1), (3, 5)])
+@pytest.mark.parametrize("input_dim", [(1,), (3,), (1, 1), (2, 1), (3, 5)])
 def test_simformer_builder_shape_and_build(
     builder,
     builder_kwargs,