Additional use of jaxtyping

[phinate]

We should probably be consistent and type the dataloaders too. Also, for cases where exact dimensions are known, we can provide those as integers, e.g.

def foo(x: Float[Array, "3 10 10 3 1"]): ...

would check the exact dimensions. See https://docs.kidger.site/jaxtyping/api/array/.