Missing validation: Silent output-shape corruption when len(Metadata.time) != batch size

[xerxeschongxian26]

Summary

When len(metadata.time) doesn't match the data batch size B, the encoder's absolute-time embedding silently broadcasts the output batch dimension to len(time) via (B, L, D) + (len(time), 1, D) in aurora/model/encoder.py, line 363. No error is raised, producing wrongly-shaped predictions with corrupted metadata.

Prevent users from introducing silent errors when constructing Batch via a validation of the requirement with appropriate error messages

Environment

• python: 3.12.12
• torch: 2.12.1
• numpy: 2.5.0
• microsoft-aurora: 1.8.0 (latest PyPI)
• platform: macOS 14.5 arm64

Verified present in 1.8.0

Reproduced on AuroraSmall and AuroraSmallPretrained

Minimal reproduction

from datetime import datetime
import torch
from aurora import AuroraSmall, Batch, Metadata

model = AuroraSmall()
model.eval()

# Data batch size B = 1, but time tuple has length 100
batch = Batch(
    surf_vars={k: torch.randn(1, 2, 17, 32) for k in ("2t", "10u", "10v", "msl")},
    static_vars={k: torch.randn(17, 32) for k in ("lsm", "z", "slt")},
    atmos_vars={k: torch.randn(1, 2, 4, 17, 32) for k in ("z", "u", "v", "t", "q")},
    metadata=Metadata(
        lat=torch.linspace(90, -90, 17),
        lon=torch.linspace(0, 360, 33)[:-1],
        time=tuple(datetime(2020, 6, 1, 12, 0) for _ in range(100)),  # single datetime element duplicated 100 times
        atmos_levels=(100, 250, 500, 850),
    ),
)

with torch.inference_mode():
    pred = model.forward(batch)

print("input  2t:", tuple(batch.surf_vars["2t"].shape))  # (1, 2, 17, 32)
print("output 2t:", tuple(pred.surf_vars["2t"].shape))   # (100, 1, 16, 32)  <-- BUG

Output:

input  2t: (1, 2, 17, 32)
output 2t: (100, 1, 16, 32)

The output batch dimension is silently inflated from 1 to 100. Sweeping len(time) over (1, 2, 3, 5, 100) yields output batch dims (1, 2, 3, 5, 100) respectively, demonstrates that it tracks len(time) exactly.

Expected behaviour

A clear error at construction or in the forward pass when len(time) != B, consistent with the existing validations for latitude and longitude in Metadata.__post_init__.

Root cause: documented contract is unenforced

aurora/model/encoder.py, Perceiver3DEncoder.forward (~L349–363):

# Absolute time embedding — built from metadata.time, so size len(time):
absolute_times_list = [t.timestamp() / 3600 for t in batch.metadata.time]
absolute_times = torch.tensor(absolute_times_list, dtype=torch.float32, device=x.device)
absolute_time_encode = absolute_time_expansion(absolute_times, self.embed_dim)
absolute_time_embed = self.absolute_time_embed(absolute_time_encode.to(dtype=dtype))
x = x + absolute_time_embed.unsqueeze(1)              # (B, L, D) + (len(time), 1, D) -> broadcasts

x carries the true batch size B (from the data tensors), but absolute_times has length len(time). When they differ, the final add broadcasts (B, L, D) + (len(time), 1, D) to (len(time), L, D). The code comment even states the intended # (B, L, D) + (B, 1, D), i.e. it assumes len(time) == B, but this is never enforced.

The ERA5 example explicitly documents the requirement:

Note that this needs to be a tuple of length one: one value for every batch element.

So the intended contract is len(time) == batch_size. Unlike the latitude/longitude monotonicity requirements (documented and enforced with ValueErrors in Metadata.__post_init__), the time length contract is documented but unguarded. Since the package ships no DataLoader and the documented workflow is manual Batch construction, validation at the Batch/Metadata boundary is the only guardrail external users have.

Secondary validation gap

• When len(atmos_levels) != data level dim, this results in a cryptic RuntimeError in normalise_atmos_var far from the root cause, rather than a clear message at construction. This can be similarly remedied via validation with an appropriate error message

from datetime import datetime
import torch
from aurora import AuroraSmall, Batch, Metadata

model = AuroraSmall()
model.eval()

batch = Batch(
    surf_vars={k: torch.randn(1, 2, 17, 32) for k in ("2t", "10u", "10v", "msl")},
    static_vars={k: torch.randn(17, 32) for k in ("lsm", "z", "slt")},
    atmos_vars={k: torch.randn(1, 2, 4, 17, 32) for k in ("z", "u", "v", "t", "q")}, # 4 atmospheric levels
    metadata=Metadata(
        lat=torch.linspace(90, -90, 17),
        lon=torch.linspace(0, 360, 33)[:-1],
        time=(datetime(2020, 6, 1, 12, 0),),
        atmos_levels=(100, 250, 500),    # 3 atmospheric levels instead of 4
    ),
)

with torch.inference_mode():
    pred = model.forward(batch)

Error:

RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 2

This cryptic error occurs during normalise() or forward pass, far from the root cause. With proper validation at construction, this would fail immediately with a clear message.

---

Proposed fix

Add an extra validation, Batch.__post_init__, similar to Metadata.__post_init__

  # `Metadata.time` must contain exactly one datetime per batch element.
  if len(self.metadata.time) != b:
      raise ValueError(
          f"`Metadata.time` has length {len(self.metadata.time)}, but the batch size "
          f"is {b}. `time` must contain exactly one entry per batch element."
      )
  # `Metadata.atmos_levels` must match the pressure-level dimension of the data.
  if len(self.metadata.atmos_levels) != c:
      raise ValueError(
          f"`Metadata.atmos_levels` has length {len(self.metadata.atmos_levels)}, but the "
          f"atmospheric variables have {c} pressure levels. These must be equal."
      )

Happy to submit a PR. The fix is small, localised, and follows the existing validation pattern.

Suggested labels: bug and good first issue, since the fix is small and self-contained
Note: This issue was drafted with AI assistance.

[wesselb]

Hey Xerxes Chong (@xerxeschongxian26)! Thanks for opening an issue and proposing a solution. The proposed solution sounds very reasonable to me. If you could open a PR and tag me, then I'd be more than happy to review. :)