Vectorize the per-sample random crop in the WB2 example

[icn5381]

Closes #29.

What changed

_subregion_crop in examples/wb2_dataloader.py no longer loops over the batch:

• the two offset vectors are drawn in one batched pull each (rng.integers(..., size=B)),
• middle dimensions fold into a flattened (B*M, LAT, LON) view, each sample's offsets repeat across its middle slices (so one spatial window per sample covers them all — same semantics as the loop),
• the crops are gathered with adjacent advanced indexing over every axis of the flattened view, then reshaped back to (batch, ..., h, w).

I first tried broadcasting (B, 1, h, 1) / (B, 1, 1, w) index blocks against trailing a[..., r, c], but NumPy appends the full broadcast block after a.shape[:-2] rather than merging the leading axes — shape came out wrong with middle dims present. The fold-and-gather form is the one I verified against a brute-force reference. Per-sample randomness is unchanged (each sample its own offset); only the RNG consumption order differs, so a given seed is not bit-identical to the old loop — as the issue anticipated, which is why the tests assert properties, not recorded values.

How I verified

New tests/test_wb2_subregion_crop.py (5 tests): full-size crop is the identity; shape/dtype preserved through middle dims; every crop is a real contiguous window of its own sample (located by search, so a transposed or hoisted axis fails); constant-filled samples prove no cross-sample bleed; same seed reproduces, different seed differs. Plus the existing test_example.py smoke tests still pass, and the full suite is green: 241 passed, 29 skipped (tests/test_zarr_indexing_parity.py errors at collection on this machine — missing optional zarr_indexing module, pre-existing and unrelated). ruff check and ruff format clean on the touched files.

[emfdavid]

Thank you for picking this up — and for writing up the broadcasting approach that didn't work. That saved me from suggesting it.

The tests are the strongest part. _find_window comparing the whole sample, middle dims included, pins "one window per sample covers all of its middle slices" — I broke that on purpose (swapping np.repeat for np.tile) and your test caught it. That's a better bar than the issue asked for. I've left one comment inline on a gap in the same file.

The main thing: worth benchmarking against the sliding-window-view form. SWV was new to me too, but the difference is bigger than I expected. Your gather indexes every axis with an array, so numpy walks it element by element and the window's contiguity buys nothing. Leaving the last two axes as slices lets it copy each window as a block:

win = np.lib.stride_tricks.sliding_window_view(a, (h, w), axis=(-2, -1))
out = win[np.arange(a.shape[0]), ..., i, j, :, :]

The ellipsis separates the advanced indices from the trailing :, :, so the batch axis lands first and the middle dims stay put — the reshape goes away too. What I measured (4 vCPU GCP box, numpy 2.4.6 / py3.12.13, float32, best of 5x200, all three forms checked against a brute-force reference first — script: https://gist.github.com/emfdavid/6daa67d2d93360bef9031b648de7ed22):

case                                              loop    this PR       swv
WB2 default    (32,64,128)    crop 16x16        0.033ms    0.048ms    0.017ms
WB2 demo-like  (32,64,128)    crop 48x48        0.053ms    0.315ms    0.040ms
wide batch     (256,64,128)   crop 16x16        0.269ms    0.343ms    0.059ms
middle dims    (64,8,121,240) crop 64x64        1.210ms   10.119ms    1.090ms

The gap widens with crop area, which is the per-element signature.

Would you run it on your machine and post what you get, with the setup line? I'd rather
have a second measurement than treat mine as settled — and if it doesn't reproduce for you, I want to know that before this merges.

To be clear: the issue told you to pick whichever form read most clearly and said nothing
about speed. That's on me — I didn't expect the performance difference.

Last thing, and it's a real question rather than a hint: I'm still tuning the PR template. Did it get in the way, or just feel like overkill for a change this size? The one part I do need is the author-attestation box — that's the whole basis of the AI policy — but the rest I'm happy to cut if it's noise.

[emfdavid]

This one can't fail on what its name promises, and it's worth fixing because the property is the whole point of the change.

I replaced your two batched draws with a single scalar offset broadcast across the batch — every sample cropped at the identical window, which is exactly the bug vectorizing this could introduce — and all five tests still passed. This one included: constant-filling sample b makes every window of it identical, so the offset is unobservable by construction. What survives is a second cross-sample-bleed check, which test_each_crop_is_a_real_contiguous_window_of_its_own_sample above already does more strictly.

The comment on lines 62-64 describes the test you meant to write ("at least two samples land on different windows"), and _find_window already returns the (i, j) it locates — it's just discarded at line 57. Something like: keep the located offsets from real (non-constant) data and assert the batch produced more than one distinct one. Fixed seed and real values, so it stays deterministic rather than probabilistic.

Small separate thing: the "sample b is filled with the constant b" comment on lines 50-52 belongs to this test, not the one it's sitting on.