Remove constraint on max_sequence_length and update related logic

olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py

@@ -840,7 +840,7 @@ def __init__(
         self,
         embedding_size: int,
         supported_modalities: list[ModalitySpec],
-        max_sequence_length: int,
+        max_sequence_length: int | None = None,
         learnable_channel_embeddings: bool = True,
         random_channel_embeddings: bool = False,
         tokenization_config: TokenizationConfig | None = None,
@@ -853,7 +853,8 @@ def __init__(
             embedding_size: Size of token embeddings
             supported_modalities: Which modalities from Modality this model
                 instantiation supports
-            max_sequence_length: Maximum sequence length
+            max_sequence_length: Size of the checkpoint-compatible temporal
+                encoding table. Longer sequences are computed on-the-fly.
             learnable_channel_embeddings: Whether to use learnable channel embeddings
             random_channel_embeddings: Initialize channel embeddings randomly (zeros if False)
             tokenization_config: Optional config for custom band groupings
@@ -870,6 +871,8 @@ def __init__(
                 f"position_encoding must be one of {PositionEncoding.values()}, "
                 f"got {position_encoding}"
             )
+        if max_sequence_length is None:
+            max_sequence_length = 12
         self.embedding_size = embedding_size
         self.supported_modalities = supported_modalities
         self.supported_modality_names = [
@@ -878,15 +881,13 @@ def __init__(
         self.tokenization_config = tokenization_config or TokenizationConfig()
         self.position_encoding = position_encoding
         self.embedding_size = embedding_size
-        self.max_sequence_length = (
-            max_sequence_length  # This max sequence length is a time dim thing
-        )
         # TODO: we need to be able to calculate the size of the param based on what types of embeddings it will get
 
         # we have 4 embeddings types (pos_in_time, pos_in_space, month, channel) so each get
         # 0.25 of the dimension
         self.embedding_dim_per_embedding_type = int(embedding_size * 0.25)
-        # Position encodings for time dimension initialized to 1D sinusoidal encodings
+        # Retain the frozen table for exact compatibility with existing
+        # checkpoints. Positions beyond it are computed on-the-fly.
         self.pos_embed = nn.Parameter(
             get_1d_sincos_pos_encoding(
                 torch.arange(max_sequence_length),
@@ -1032,10 +1033,15 @@ def _apply_encodings_per_modality(
             # Slot-index temporal encoding (additive). Skipped when 3D RoPE
             # handles temporal position rotationally inside attention.
             if not PositionEncoding.is_3d_rope(self.position_encoding):
-                time_embed = repeat(
-                    self.pos_embed[:t], f"t d -> {ein_string}", **ein_dict
-                )
-                modality_embed[..., n : n * 2] += time_embed.to(device)
+                if t <= self.pos_embed.shape[0]:
+                    pos_embed = self.pos_embed[:t].to(device)
+                else:
+                    pos_embed = get_1d_sincos_pos_encoding(
+                        torch.arange(t, device=device),
+                        self.embedding_dim_per_embedding_type,
+                    )
+                time_embed = repeat(pos_embed, f"t d -> {ein_string}", **ein_dict)
+                modality_embed[..., n : n * 2] += time_embed
 
             # Month encodings stay additive in all modes (calendar/seasonal
             # signal is orthogonal to slot-index).

olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py

@@ -4,7 +4,7 @@
 """
 
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
 
 # The highest resolution that we are working at.
 # Everything else is a factor (which is a power of 2) coarser than this resolution.
@@ -81,7 +81,7 @@ def get_expected_image_size(self, modality_resolution_factor: int) -> int:
         return IMAGE_TILE_SIZE // (self.resolution_factor // modality_resolution_factor)
 
 
-class TimeSpan(str, Enum):
+class TimeSpan(StrEnum):
     """Enum to distinguish data that is valid for different time ranges."""
 
     # Only one data point (not time series).

tests/test_flexi_vit.py

@@ -1,7 +1,58 @@
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import Encoder
+import torch
+
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.encodings import (
+    get_1d_sincos_pos_encoding,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import (
+    CompositeEncodings,
+    Encoder,
+)
 from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality
 
 
+def test_composite_encodings_compute_temporal_positions_on_the_fly() -> None:
+    """CompositeEncodings supports sequences longer than deprecated config value."""
+    encodings = CompositeEncodings(
+        embedding_size=16,
+        supported_modalities=[Modality.SENTINEL2_L2A],
+        max_sequence_length=3,
+    )
+    assert encodings.pos_embed.shape == (3, 4)
+
+    tokens = torch.zeros((1, 2, 2, 5, Modality.SENTINEL2_L2A.num_band_sets, 16))
+    timestamps = torch.tensor(
+        [[[1, 0, 2020], [2, 1, 2020], [3, 2, 2020], [4, 3, 2020], [5, 4, 2020]]]
+    )
+
+    output = encodings(
+        {Modality.SENTINEL2_L2A.name: tokens},
+        timestamps=timestamps,
+        patch_size=4,
+    )
+
+    encoded = output[Modality.SENTINEL2_L2A.name]
+    assert encoded.shape == tokens.shape
+    assert torch.isfinite(encoded).all()
+
+    expected_time = get_1d_sincos_pos_encoding(torch.arange(5), 4)
+    actual_time = encoded[0, 0, 0, :, 0, 4:8]
+    assert torch.allclose(actual_time, expected_time, atol=1e-5)
+
+
+def test_composite_encodings_loads_temporal_position_table() -> None:
+    """Strict state dict loading restores the frozen temporal table."""
+    encodings = CompositeEncodings(
+        embedding_size=16,
+        supported_modalities=[Modality.SENTINEL2_L2A],
+        max_sequence_length=3,
+    )
+    state_dict = encodings.state_dict()
+    state_dict["pos_embed"] = torch.zeros((3, 4))
+
+    encodings.load_state_dict(state_dict, strict=True)
+    assert torch.equal(encodings.pos_embed, torch.zeros((3, 4)))
+
+
 def test_band_dropout_disabled_by_default() -> None:
     """Test that Encoder leaves band dropout disabled at construction."""
     encoder = Encoder(