address review comments

Author: mnabian

physicsnemo/nn/module/embedding_layers.py

@@ -299,9 +299,10 @@ class FourierPositionalEmbedding(Module):
     freq_scale : float, optional, default=2.0
         Geometric ratio between consecutive band frequencies for the
         generated schedule.
-    freqs : torch.Tensor or sequence of float, optional, default=None
-        Explicit 1-D frequency schedule. Overrides ``num_bands``,
-        ``base_freq`` and ``freq_scale`` when provided.
+    freqs : torch.Tensor, optional, default=None
+        Explicit 1-D frequency schedule of shape :math:`(F,)` with
+        :math:`F \geq 1`. Overrides ``num_bands``, ``base_freq`` and
+        ``freq_scale`` when provided. A non-1-D ``freqs`` raises ``ValueError``.
 
     Forward
     -------
@@ -334,7 +335,7 @@ def __init__(
         include_input: bool = True,
         base_freq: float = math.pi,
         freq_scale: float = 2.0,
-        freqs: Tensor | None = None,
+        freqs: Float[Tensor, "num_freqs"] | None = None,  # noqa: F821
     ):
         super().__init__()
         if in_dim < 1:
@@ -346,9 +347,12 @@ def __init__(
                 num_bands, dtype=torch.float32
             )
         else:
-            freqs = torch.as_tensor(freqs, dtype=torch.float32).flatten()
-            if freqs.numel() < 1:
-                raise ValueError("freqs must contain at least one frequency.")
+            freqs = freqs.to(torch.float32)
+            if freqs.ndim != 1 or freqs.numel() < 1:
+                raise ValueError(
+                    "freqs must be a 1-D tensor of shape (F,) with F >= 1, "
+                    f"got shape {tuple(freqs.shape)}."
+                )
         self.in_dim = int(in_dim)
         self.include_input = bool(include_input)
         # Persistent so an explicitly supplied ``freqs`` schedule survives a
@@ -367,7 +371,9 @@ def out_dim(self) -> int:
         base = self.in_dim if self.include_input else 0
         return base + 2 * self.in_dim * self.num_bands
 
-    def forward(self, x: Float[Tensor, "... in_dim"]) -> Float[Tensor, "... out_dim"]:
+    def forward(
+        self, x: Float[Tensor, "*dims in_dim"]
+    ) -> Float[Tensor, "*dims out_dim"]:
         r"""Encode coordinates ``x``; see the class docstring for shapes."""
         # Skip validation when running under torch.compile (MOD-005).
         if not torch.compiler.is_compiling():

test/nn/module/data/fourier_positional_embedding_in3_nb4_b2x4.pth

@@ -425,6 +425,13 @@ def test_deterministic(self, device):
             False,
             2 * 2 * 6,
         ),  # non-defaults
+        (
+            {"in_dim": 3, "freqs": torch.tensor([1.0, 2.0, 4.0])},
+            3,
+            3,
+            True,
+            3 + 2 * 3 * 3,
+        ),  # explicit freqs (num_bands inferred from the schedule length)
     ],
 )
 def test_fourier_positional_embedding_constructor_attrs(
@@ -435,79 +442,60 @@ def test_fourier_positional_embedding_constructor_attrs(
     assert emb.num_bands == exp_num_bands
     assert emb.include_input == exp_include_input
     assert emb.out_dim == exp_out_dim
-
-
-def test_fourier_positional_embedding_out_dim_and_shape(device):
-    emb = FourierPositionalEmbedding(in_dim=3, num_bands=4).to(device)
-    assert emb.num_bands == 4
-    assert emb.out_dim == 3 + 2 * 3 * 4  # 27
-    y = emb(torch.randn(5, 3, device=device))
-    assert y.shape == (5, 27)
     # No learnable parameters.
     assert sum(p.numel() for p in emb.parameters()) == 0
 
 
-def test_fourier_positional_embedding_no_include_input(device):
-    emb = FourierPositionalEmbedding(in_dim=2, num_bands=3, include_input=False).to(
-        device
-    )
-    assert emb.out_dim == 2 * 2 * 3  # 12
-    assert emb(torch.zeros(4, 2, device=device)).shape == (4, 12)
-
-
-def test_fourier_positional_embedding_leading_dims(device):
-    emb = FourierPositionalEmbedding(in_dim=3, num_bands=2).to(device)
-    for shape in [(3,), (7, 3), (2, 7, 3)]:
-        out = emb(torch.randn(*shape, device=device))
-        assert out.shape == (*shape[:-1], emb.out_dim)
-
-
-def test_fourier_positional_embedding_values(device):
-    # Single coord, single band at base_freq=pi -> [sin(pi*x), cos(pi*x)].
-    emb = FourierPositionalEmbedding(in_dim=1, num_bands=1, include_input=False).to(
-        device
-    )
-    x = torch.tensor([[0.5]], device=device)
-    f = math.pi
-    torch.testing.assert_close(
-        emb(x),
-        torch.tensor([[math.sin(f * 0.5), math.cos(f * 0.5)]], device=device),
-    )
-
-
-def test_fourier_positional_embedding_explicit_freqs(device):
-    emb = FourierPositionalEmbedding(
-        in_dim=3, freqs=torch.tensor([1.0, 2.0, 4.0]), include_input=True
-    ).to(device)
-    assert emb.num_bands == 3
-    assert emb.out_dim == 3 + 2 * 3 * 3
-    assert emb(torch.randn(8, 3, device=device)).shape == (8, emb.out_dim)
-
-
-def test_fourier_positional_embedding_axis_major_layout(device):
-    # With include_input=False the layout is axis-major: for each axis, the
-    # num_bands sines followed by the num_bands cosines.
+@pytest.mark.parametrize(
+    "in_dim, freqs, include_input, x, expected",
+    [
+        # include_input=False, axis-major layout: per axis, sines then cosines.
+        (
+            2,
+            [1.0, 2.0],
+            False,
+            [[0.3, 0.7]],
+            [
+                [
+                    math.sin(1.0 * 0.3),
+                    math.sin(2.0 * 0.3),
+                    math.cos(1.0 * 0.3),
+                    math.cos(2.0 * 0.3),
+                    math.sin(1.0 * 0.7),
+                    math.sin(2.0 * 0.7),
+                    math.cos(1.0 * 0.7),
+                    math.cos(2.0 * 0.7),
+                ]
+            ],
+        ),
+        # Single coordinate and band.
+        (
+            1,
+            [math.pi],
+            False,
+            [[0.5]],
+            [[math.sin(math.pi * 0.5), math.cos(math.pi * 0.5)]],
+        ),
+        # include_input=True prepends the raw coordinate.
+        (
+            1,
+            [1.0],
+            True,
+            [[0.5]],
+            [[0.5, math.sin(0.5), math.cos(0.5)]],
+        ),
+    ],
+)
+def test_fourier_positional_embedding_forward_values(
+    device, in_dim, freqs, include_input, x, expected
+):
+    # Known-reference forward values across configs (layout, single band,
+    # and include_input prepend).
     emb = FourierPositionalEmbedding(
-        in_dim=2, num_bands=2, include_input=False, freqs=torch.tensor([1.0, 2.0])
+        in_dim=in_dim, freqs=torch.tensor(freqs), include_input=include_input
     ).to(device)
-    x = torch.tensor([[0.3, 0.7]], device=device)
-    out = emb(x)
-    expected = torch.tensor(
-        [
-            [
-                math.sin(1.0 * 0.3),
-                math.sin(2.0 * 0.3),
-                math.cos(1.0 * 0.3),
-                math.cos(2.0 * 0.3),
-                math.sin(1.0 * 0.7),
-                math.sin(2.0 * 0.7),
-                math.cos(1.0 * 0.7),
-                math.cos(2.0 * 0.7),
-            ]
-        ],
-        device=device,
-    )
-    torch.testing.assert_close(out, expected)
+    out = emb(torch.tensor(x, device=device))
+    torch.testing.assert_close(out, torch.tensor(expected, device=device))
 
 
 def test_fourier_positional_embedding_validation(device):
@@ -518,6 +506,9 @@ def test_fourier_positional_embedding_validation(device):
         FourierPositionalEmbedding(in_dim=0)
     with pytest.raises(ValueError):
         FourierPositionalEmbedding(in_dim=3, num_bands=0)
+    # Explicit freqs must be 1-D of shape (F,).
+    with pytest.raises(ValueError):
+        FourierPositionalEmbedding(in_dim=3, freqs=torch.ones(2, 3))
 
 
 def test_fourier_positional_embedding_state_dict_roundtrip(device):
@@ -540,12 +531,13 @@ def test_fourier_positional_embedding_forward_accuracy(device):
     # MOD-008b: compare the forward output against committed reference data.
     model = FourierPositionalEmbedding(in_dim=3, num_bands=4).to(device)
     model.eval()
-    # Deterministic, reproducible input (the layer has no random parameters).
-    x = torch.linspace(-1.0, 1.0, steps=24, device=device).reshape(8, 3)
+    # Deterministic, reproducible input; a 3-D shape also exercises arbitrary
+    # leading (batch) dimensions against the reference.
+    x = torch.linspace(-1.0, 1.0, steps=2 * 4 * 3, device=device).reshape(2, 4, 3)
     assert validate_forward_accuracy(
         model,
         (x,),
-        file_name="nn/module/data/fourier_positional_embedding_in3_nb4_bs8.pth",
+        file_name="nn/module/data/fourier_positional_embedding_in3_nb4_b2x4.pth",
         rtol=1e-4,
         atol=1e-4,
     )