Data and datasets

Data records

data Name: followed by an indented block of fields declares a data record: the fields of one example of a dataset.

## One labelled image.
data Digit:
    img:   [D]
    label: [] Index[C]

Each field is a name and a tensor type, and the record and its fields can have doc comments. A field’s shape is made only of numbers and dims, so every example has the same shapes and examples can be stacked into batches: a generic name is E0301, and ?k is E1302. A field can have any dtype and any domain. It can’t be a host scalar, a type without a shape, a model, a data or logs record, a dataset, a gradient, optimizer state, a Schedule, a Mode or a LogSink, and it can’t be named with a Python keyword or a name that starts with _ (E1302). A record’s fields have different names (E0302).

Values of different lengths, such as sentences, aren’t supported yet: every example of a record has the same shapes, and an example whose field has another shape is an error when it’s loaded. Pad them to a dim in the dataset and add a Mask field, as in tokens: [T] Index[V] and mask: [T] Mask. Padding varying axes for you is planned (FAST-3 in the plan).

Batches

[b] Digit is a batch of b examples, in which every field has the batch’s axes first: batch.img is [b D] and batch.label is [b] Index[C]. Digit alone is one example, whose fields have their declared types. A field the record doesn’t declare is E1301.

A batch type can be a parameter or return type, and in a parameter type, as with a tensor, a new name such as b is a generic axis that the arguments determine (Generic calls). Only a batch of the same record with the same number of batch axes fits; another record, a dataset or a tensor is E1303. A batch isn’t a tensor: arithmetic on one is E0412. A record’s name isn’t a value or a function (E0303), and a record’s name with a size, an over or a dtype, as in Digit[3], is E0303 too. as can’t give a value a batch type (E0303).

Datasets

Dataset[Digit] is a dataset of Digit examples. A dataset comes from Python: a parameter of that type takes any map-style dataset, an object with __len__ and __getitem__ such as a torch.utils.data.Dataset, whose items have the record’s fields, as a dict (other keys are ignored) or as a tuple or list in field order; for a record with one field, an item may be its value. In a notebook, x as Dataset[Digit] gives a Python value that type. Only a dataset of the same record fits a Dataset[Digit] parameter or return type (E1303). Dataset with anything but a record’s name in its brackets, or with a shape, is E0303.

Datasets are used with two pipe stages, which are built-in functions:

seed is given by name and n by position or as n= (E0406; E0502 or E0805 for something other than an int). Either stage on something that isn’t a dataset, shuffle after batches, batches twice, and a batch size known to be less than 1 are E1304. The order shuffle gives comes from a generator of its own, seeded from s, so it doesn’t depend on PyTorch’s global random state, and the same seed always gives the same order.

Looping over batches

for batch in train |> shuffle(seed=epoch) |> batches(64): runs its block once per batch. All batches but the last have n examples; the last has the examples that are left, so 100 examples in batches of 64 give a batch of 64 and one of 36. So the batch axis’s size isn’t known statically: the loop variable is [b] Digit, where b is an axis whose size is unknown, a new one for each loop, shown as ?b where it’s compared with other sizes. Inside the loop, b names that axis, as in batch.img as [b D], and is its size as a value, unless b already names something. Passing the batch to a function with a [b] Digit parameter makes the function’s b that size. The loader pads the last batch to the full size with masked-out examples, so that every batch has the same shape, when every operation across the batch has a masked version; b is then still the number of real examples where it’s used as a value (Fixed shapes and padding).

Loading

When a loop starts, the loader loads the examples in the loop’s order and checks each against the record as it’s loaded: the item must have every field, each field must be a tensor, a number, or anything torch.as_tensor takes, with the field’s shape exactly, and its values must be of the field’s kind. A float field takes any numbers, an integer field integers or booleans, and a bool field booleans; each is then converted to the field’s dtype. With value checks on (Runtime checks), each field is also checked against its domain. An item that doesn’t fit raises shapescript.ShapeError, or shapescript.DomainError for a domain, naming the .ss location of the field, the field and the item’s index in the dataset:

ShapeError: mnist.ss:8:5: field `img` of item 37 should be [D] = (784,), got shape (783,)

The examples of a batch are stacked field by field, and the batch is copied to the declared device. From Python, a batch is a shapescript.rt.Batch, whose fields are attributes. A function with a batch parameter also takes a dict with the record’s fields from Python, and checks each field as it checks a tensor argument.

Worker processes. Batches are built by worker processes, 4 unless the machine has fewer CPUs; shapescript.loader_workers(n) changes that, and 0 loads them in the calling process. The workers are started once per dataset parameter and kept for the loops over it. They are always started fresh, as with multiprocessing’s “spawn” start method, whatever the start method is, and never import the program’s main module: a worker unpickles the dataset, which imports the dataset’s own module, and nothing else. So the top level of a calling script runs once, and needs no if __name__ == "__main__":, even where Python’s default start method would import it again in each worker (“spawn” on macOS and Windows, and “forkserver” on Linux from Python 3.14). A dataset that refers to something defined in the main script, or that can’t be pickled, is loaded in the calling process instead, with a RuntimeWarning saying why. A module built with --standalone always loads in the calling process, since workers can’t import it.

Datasets kept on the device. A PyTorch torch.utils.data.TensorDataset with a tensor per field, in field order, whose item i is row i of each, isn’t loaded batch by batch. When a loop over it starts on the CPU, or on a GPU where it takes at most a quarter of the free memory, it’s checked once, as a whole, against the record, converted to the fields’ dtypes and copied to the device, and each batch is gathered there with an index operation: no loader workers, no copy from the CPU and no check per batch. It’s copied and checked again only if its tensors change (an in-place change to one counts). A subclass that changes __getitem__ or __len__, or adds __getitems__, as an augmentation does, is loaded as any dataset is; so is one with an item that doesn’t fit, which raises the error as the loader would. On several devices each worker keeps a copy of its own.