ArrayDataset

jax_dataloader.datasets.ArrayDataset

Bases: Dataset

Dataset wrapping numpy arrays.

Source code in jax_dataloader/datasets.py

class ArrayDataset(Dataset):
    """Dataset wrapping numpy arrays."""

    def __init__(
        self, 
        *arrays: jax.Array, # Numpy array with same first dimension
        asnumpy: bool = True, # Store arrays as numpy arrays if True; otherwise store as array type of *arrays
    ):
        assert all(arrays[0].shape[0] == arr.shape[0] for arr in arrays), \
            "All arrays must have the same dimension."
        self.arrays = tuple(arrays)
        if asnumpy:
            self.asnumpy()            

    def asnumpy(self):
        """Convert all arrays to numpy arrays."""
        self.arrays = tuple(asnumpy(arr) for arr in self.arrays)

    def __len__(self):
        return self.arrays[0].shape[0]

    def __getitem__(self, index):
        return jax.tree_util.tree_map(lambda x: x[index], self.arrays)

Attributes

arrays = tuple(arrays) instance-attribute

Functions

__getitem__(index)

Source code in jax_dataloader/datasets.py

def __getitem__(self, index):
    return jax.tree_util.tree_map(lambda x: x[index], self.arrays)

__init__(*arrays: jax.Array, asnumpy: bool = True)

Source code in jax_dataloader/datasets.py

def __init__(
    self, 
    *arrays: jax.Array, # Numpy array with same first dimension
    asnumpy: bool = True, # Store arrays as numpy arrays if True; otherwise store as array type of *arrays
):
    assert all(arrays[0].shape[0] == arr.shape[0] for arr in arrays), \
        "All arrays must have the same dimension."
    self.arrays = tuple(arrays)
    if asnumpy:
        self.asnumpy()            

__len__()

Source code in jax_dataloader/datasets.py

def __len__(self):
    return self.arrays[0].shape[0]

asnumpy()

Convert all arrays to numpy arrays.

Source code in jax_dataloader/datasets.py

def asnumpy(self):
    """Convert all arrays to numpy arrays."""
    self.arrays = tuple(asnumpy(arr) for arr in self.arrays)