Validate that metadata.time matches the batch size

[priyDe (priyald)]

Fixes #188.

What's changed since the issue was filed

The issue reports silent corruption on 1.8.0: the output batch dimension is inflated to
len(metadata.time) with no error. On current main the symptom is different —
len(metadata.time) == 1 still works, and anything larger now fails with

IndexError: index 1 is out of bounds for dimension 0 with size 1

from lead_times[i] in Perceiver3DDecoder.forward. The root cause is unchanged: the encoder
still adds the absolute time embedding as

x = x + absolute_time_embed.unsqueeze(1)  # (B, L, D) + (B, 1, D)

and that comment states an invariant which nothing enforces. So the failure has moved from
"quietly wrong" to "loud but misleading", and in the len(metadata.time) < B direction it is
still quietly wrong.

The same bug on the pressure-level axis

While adding the check I looked for the same pattern elsewhere, and Aurora3DEncoder.forward has
it:

atmos_levels_embed = self.atmos_levels_embed(atmos_levels_encode)[None, :, None, :]
x_atmos = x_atmos + atmos_levels_embed  # (B, C_A, L, D)

x_atmos is (B, C, L, D) and the encoding is (1, len(atmos_levels), 1, D), so a disagreement
between len(metadata.atmos_levels) and the level dimension of the atmospheric variables
broadcasts exactly as the time embedding does. Running it on AuroraSmall:

Input	Result on main
1 level of data, 4 atmos_levels	runs, returns 4 pressure levels — three fabricated
4 levels of data, 1 atmos_level	runs, returns 1 pressure level — three silently dropped
4 levels of data, 2 atmos_levels	RuntimeError from the addition (loud, so fine)

The two silent cases are the same class of bug as #188, so I've handled both here rather than
opening a near-identical second PR. Happy to split it out if you'd prefer to review them
separately.

The fix

Batch.__post_init__ now checks that

• all surface-level and atmospheric variables agree on the batch dimension,
• metadata.time gives exactly one time per batch element,
• all atmospheric variables agree on the pressure-level dimension, and
• metadata.atmos_levels gives exactly one level per level of those variables,

raising a ValueError at construction rather than deferring the failure. This matches the
documented contract — "time: For every batch element, the time" — and is consistent with the
existing latitude and longitude validation in Metadata.__post_init__. The two checks share a
small _consistent_size helper so the logic exists once.

One implementation note: the level dimension is counted from the end (dim=-3). Perceiver3DDecoder
constructs a Batch whose atmospheric variables are (B, C, H, W), before Aurora.forward inserts
the history dimension, so a positive index would read the height as a level count there.

One test needed changing, and it's worth flagging

test_aurora_small duplicates the data to batch size two but leaves metadata.time at length
one:

batch = dataclasses.replace(
    batch,
    surf_vars={k: v.repeat(2, 1, 1, 1) for k, v in batch.surf_vars.items()},
    atmos_vars={k: v.repeat(2, 1, 1, 1, 1) for k, v in batch.atmos_vars.items()},
)

It passes on main because the broadcast happens to do the right thing in that direction, but
it does contradict the documented contract. I've repeated the time along with the data. Happy to
take a different approach if you'd rather the len(metadata.time) == 1 case stay supported as an
explicit broadcast — that would be a small change to the check.

Tests

Added to tests/test_batch.py:

• matching times and batch size is accepted, for batch sizes 1, 2 and 3;
• time mismatches raise, covering both directions — (B=1, times=2), (B=1, times=100),
  (B=2, times=1), (B=3, times=2);
• variables disagreeing with each other on the batch dimension raise;
• matching levels and atmos_levels is accepted, for 1, 2 and 4 levels;
• level mismatches raise, again both directions — (C=1, levels=4), (C=4, levels=1),
  (C=4, levels=2), (C=2, levels=4);
• atmospheric variables disagreeing with each other on the level dimension raise;
• to, crop and regrid still produce valid batches.

tests/test_batch.py and tests/test_headers.py pass (73 tests). The full suite is
117 passed, 1 failed, the failure being the pre-existing test_aurora_small described below.
ruff check and ruff format --check are clean.

Note on test_aurora_small

It fails for me at 0.04224444 / 276.22, marginally over the 1e-4 tolerance. This is
pre-existing and unrelated: I get the identical value on unmodified main, and it looks like
the machine-dependent numerics discussed in #169. Flagging it so it isn't mistaken for a
regression here.

[priyDe (priyald)]

Wessel (@wesselb) tagging you as you offered on #188. Happy to change the approach if you'd rather keep len(metadata.time) == 1 supported as an explicit broadcast, that would be a small edit to the check.

[priyDe (priyald)]

@microsoft-github-policy-service agree

[priyDe (priyald)]

Pushed a second commit (bc00cec) while waiting on review, and updated the description to match.

Looking for the same pattern elsewhere turned up one more instance: the pressure encoding in Aurora3DEncoder.forward is added as (B, C, L, D) + (1, len(atmos_levels), 1, D), so a disagreement between len(metadata.atmos_levels) and the level dimension broadcasts the same way the time embedding does. On AuroraSmall, one level of data with four atmos_levels returns four levels with three fabricated, and four levels of data with one atmos_level returns one level, silently dropping three. Both are now caught at construction, and the two checks share a small helper.

It seemed better to fix the same bug on both axes in one place than to open a near-identical second PR, but I am happy to split it if you would rather review them separately.

[priyDe (priyald)]

Closing this in favour of #196, which fixes the same issue and was opened six weeks earlier. I should have found it before opening this one — it does not appear on the #188 thread, but that is my miss, not an excuse. Apologies for the extra tab, Wessel (@wesselb).

For what it is worth, the two arrived at the same design independently, including using shape[-3] for the pressure-level check so that the decoder's intermediate (B, C, H, W) batch validates correctly. That seems like good evidence the approach is the right one.

Two small gaps I ran into that #196 does not cover, if they are wanted once it lands:

• it reads only the first surface and atmospheric variable via next(iter(...)), so variables that disagree with each other on the batch or level dimension are not caught;
• an empty surf_vars or atmos_vars raises StopIteration rather than a clear error.

Happy to open a small follow-up for those after #196 is merged, or to leave it be if you would rather keep the validation minimal. Either is fine by me.