ArrayDataset

class braintools.trainer.ArrayDataset(*arrays)

Dataset wrapping arrays.

Parameters:

*arrays (Array-like) – Arrays to wrap. All arrays must have the same length (first dimension).

Examples

>>> X = jnp.ones((100, 10))
>>> y = jnp.zeros((100,))
>>> dataset = ArrayDataset(X, y)
>>> len(dataset)
100
>>> x_sample, y_sample = dataset[0]