ArrayLoader

ArrayLoader(dataset, rng_key, *[, ...])

Data loader for batching and padding arrays.

Methods

ArrayLoader.pad_array(x, n_pad[, axis])

Pad an array to ensure compatibility with sharding.