Typing Extensions
About
The d9d.core.types package holds common type aliases used across d9d, such as PyTree and TensorTree. The d9d.core.protocol package defines protocols for standard PyTorch components that the training loop uses, such as the optimizer, the learning rate scheduler and the data loader.
API Reference
d9d.core.types
Common type definitions used across d9d.
CollateFn = Callable[[Sequence[TDataTree]], TDataTree]
module-attribute
Type alias for a function that collates a sequence of samples into a batch.
The function receives a sequence of samples (PyTrees) and stacks or merges them into one batched PyTree.
MicrobatchPack = Sequence[TDataTree]
module-attribute
Type alias for one step's worth of data: a sequence of microbatches.
PyTree = TLeaf | list['PyTree[TLeaf]'] | dict[str, 'PyTree[TLeaf]'] | tuple['PyTree[TLeaf]', ...]
module-attribute
Type alias for a recursive tree of data.
The tree nests standard Python containers (dicts, lists, tuples) to any depth. Its leaves have type TLeaf.
It describes nested state dicts and arguments of functions that traverse them recursively, similar to
torch.utils._pytree.
ScalarTree = PyTree[str | float | int | bool]
module-attribute
Type alias for a PyTree whose leaves are Python scalars (str, float, int, bool).
TensorTree = PyTree[torch.Tensor]
module-attribute
Type alias for a PyTree whose leaves are tensors.
d9d.core.protocol
Protocol definitions for standard PyTorch objects.
DataLoaderProtocol
Bases: Protocol
Protocol for a sized, stateful stream of single microbatches.
A conforming loader yields one collated microbatch at a time and reports its length in microbatches.
It supports checkpointing through the Stateful interface (state_dict and load_state_dict).
A torchdata StatefulDataLoader conforms to it.
__iter__()
__len__()
Returns the number of microbatches this loader yields.
Returns:
| Type | Description |
|---|---|
int
|
The number of microbatches. |
load_state_dict(state_dict)
LRSchedulerProtocol
Bases: Protocol
Protocol for a learning rate scheduler.
A conforming scheduler supports stepping and checkpointing through the Stateful interface
(state_dict and load_state_dict).
load_state_dict(state_dict)
state_dict()
step()
Performs a single learning rate scheduling step.
MicrobatchPackStream
Bases: Protocol
Protocol for a stateful stream of microbatch packs that the loop drives.
One pack holds exactly one step's worth of microbatches. The stream supports checkpointing through the
Stateful interface (state_dict and load_state_dict). It is the only checkpoint boundary of the
data pipeline.
The stream yields CPU tensors, which can be in pinned memory. The loop moves each pack to the device. When prefetching, the loop iterates the stream on a background thread.
total_steps
property
The number of steps (packs) this stream yields, or None if it is not known ahead of time.
A streaming or data-dependent source returns None. In that case, the job duration must come from
JobScheduleConfig.
__iter__()
Returns an iterator over microbatch packs.
Returns:
| Type | Description |
|---|---|
Iterator[MicrobatchPack]
|
An iterator yielding one microbatch pack (one step's worth of microbatches) at a time. |
load_state_dict(state_dict)
OptimizerProtocol
Bases: Protocol
Protocol for a standard PyTorch optimizer.
A conforming optimizer supports stepping, zeroing gradients and checkpointing through the Stateful
interface (state_dict and load_state_dict).
load_state_dict(state_dict)
state_dict()
step()
Performs a single optimization step.
zero_grad()
Sets the gradients of all optimized tensors to zero.