From e7812f801b0109c6d8d4f65a1e987e51cd56e0e7 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Wed, 26 Aug 2026 12:15:26 -0500 Subject: [PATCH 1/2] Handle two-dimensional time embeddings Avoid a zero frequency denominator for the single-frequency case while preserving existing values for larger dimensions. --- pyhealth/models/unified_embedding.py | 4 +++- tests/core/test_unified_embedding.py | 29 ++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 tests/core/test_unified_embedding.py diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py index 014326b41..3ec00576e 100644 --- a/pyhealth/models/unified_embedding.py +++ b/pyhealth/models/unified_embedding.py @@ -57,7 +57,9 @@ def __init__(self, dim: int, max_hours: float = 720.0): self.max_hours = max_hours half = dim // 2 freqs = torch.exp( - -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1) + -math.log(10000.0) + * torch.arange(half, dtype=torch.float32) + / max(half - 1, 1) ) self.register_buffer("freqs", freqs) # (dim//2,) diff --git a/tests/core/test_unified_embedding.py b/tests/core/test_unified_embedding.py new file mode 100644 index 000000000..1b34c0d5a --- /dev/null +++ b/tests/core/test_unified_embedding.py @@ -0,0 +1,29 @@ +import math +import unittest + +import torch + +from pyhealth.models import SinusoidalTimeEmbedding + + +class TestSinusoidalTimeEmbedding(unittest.TestCase): + def test_dim_two_is_finite(self): + embedding = SinusoidalTimeEmbedding(dim=2, max_hours=24.0) + + result = embedding(torch.tensor([0.0, 6.0, 24.0])) + + self.assertEqual(result.shape, (3, 2)) + self.assertTrue(torch.isfinite(result).all()) + + def test_standard_dimension_values_are_unchanged(self): + embedding = SinusoidalTimeEmbedding(dim=6, max_hours=24.0) + time = torch.tensor([6.0]) + frequencies = torch.tensor([1.0, 0.01, 0.0001]) + arguments = time.unsqueeze(-1) / 24.0 * 2 * math.pi * frequencies + expected = torch.cat([arguments.sin(), arguments.cos()], dim=-1) + + self.assertTrue(torch.allclose(embedding(time), expected)) + + +if __name__ == "__main__": + unittest.main() From 210a3ef6ec86ffdaca0d810ce46f39db44e28a3d Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 27 Aug 2026 17:52:53 -0500 Subject: [PATCH 2/2] Document sinusoidal time embedding usage --- pyhealth/models/unified_embedding.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py index 3ec00576e..7619d2de5 100644 --- a/pyhealth/models/unified_embedding.py +++ b/pyhealth/models/unified_embedding.py @@ -40,6 +40,14 @@ class SinusoidalTimeEmbedding(nn.Module): Identical in spirit to the positional encoding in "Attention is All You Need" but operating on real-valued timestamps rather than integer positions. + Examples: + >>> embedding = SinusoidalTimeEmbedding(dim=2, max_hours=24.0) + >>> output = embedding(torch.tensor([0.0, 6.0])) + >>> output.shape + torch.Size([2, 2]) + >>> torch.isfinite(output).all().item() + True + Args: dim: Output embedding dimension (must be even). max_hours: Maximum expected time value in hours. Values are normalised @@ -57,9 +65,7 @@ def __init__(self, dim: int, max_hours: float = 720.0): self.max_hours = max_hours half = dim // 2 freqs = torch.exp( - -math.log(10000.0) - * torch.arange(half, dtype=torch.float32) - / max(half - 1, 1) + -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1) ) self.register_buffer("freqs", freqs) # (dim//2,)