Favorite way to go from netCDF (&xarray) to torch/TF/Jax et al

[emfdavid]

The IO is fast now but there is still no standard tooling for chunked zarr compatible split/shuffle aware batching.

I put some time (and tokens) into building GitHub - emfdavid/insitubatch: Train in place on n-dimensional cloud tensors. · GitHub this week.

It is one Data Set abstraction compatible with Torch, Jax and TF.
It is batteries included with tunable caching, pre-fechting and max-inflight.
No heavy worker processes makes it much more efficient - lean memory and cpu (fetch once!)
It is compatible with any file/blob storage supported by obstore (aws, gcp, azure etc) and includes examples with IceChunk/ArrayLake
Check out the benchmarks and working ERA5 CNN forecast examples

from insitubatch import open_geometries, split_by_chunk
from insitubatch.source import InSituDataset

url = "file:///data/era5.zarr"           # or "s3://bucket/era5.zarr" — same code
geoms = open_geometries(url)             # {var: ArrayGeometry} from zarr metadata
manifest = split_by_chunk(geoms["t2m"], fractions=(0.8, 0.1, 0.1))

ds = InSituDataset(url, manifest, batch_size=32, block_chunks=16)

for epoch in range(n_epochs):
  ds.set_epoch(epoch)
  for batch in ds.train:               # numpy Batch: {var: np.ndarray} + sample_indices
  ...
  for batch in ds.val:                 # deterministic; shares the pool with train
  ...