DictDataset

class braintools.trainer.DictDataset(data)

Dataset wrapping a dictionary of arrays.

Parameters:

data (Dict[str, Any]) – Dictionary mapping keys to arrays.

Examples

>>> data = {'x': jnp.ones((100, 10)), 'y': jnp.zeros((100,))}
>>> dataset = DictDataset(data)
>>> len(dataset)
100
>>> sample = dataset[0]
>>> sample['x'].shape
(10,)