omniloader.transforms.base

Transform base classes and shared helpers.

A Transform operates on an already-unified sample dict (the output of SampleUnifier) and is modality-agnostic — it reasons only about vectors, sequences and their <name>_mask companions. Each transform is:

• seedable — an optional torch.Generator makes every draw reproducible;

• mask-aware — values are only altered where the mask is True (padding and placeholders are left untouched);

• train/eval-aware — augmentations set train_only = True so they are skipped during evaluation, while normalization runs in both phases.

Attributes

Sample

Classes

Transform	Base class for per-sample transforms.
Compose	Apply a list of transforms in order.

Functions

broadcast_mask(→ torch.Tensor)	Broadcast a validity mask to the value's shape for torch.where.
resolve_keys(→ list[str])	Return explicit keys or all feature keys from schema.
resolve_sequence_keys(→ list[str])	Return explicit keys or all sequence feature keys from schema.

Module Contents

omniloader.transforms.base.Sample

class omniloader.transforms.base.Transform

Bases: abc.ABC

Base class for per-sample transforms.

train_only

When True the transform is skipped unless training is passed as True to __call__().

__call__(sample: Sample, *, training: bool = True, generator: torch.Generator | None = None) → Sample

Apply the transform, honouring the train/eval gate.

Parameters:

• sample – The unified sample dict to transform.

• training – Whether the pipeline is in training mode.

• generator – Optional RNG for reproducible randomness.

Returns:

The transformed sample (mutated in place and returned).

abstractmethod apply(sample: Sample, generator: torch.Generator | None) → Sample

Transform sample unconditionally. Subclasses implement this.

class omniloader.transforms.base.Compose(transforms: collections.abc.Sequence[Transform])

Bases: Transform

Apply a list of transforms in order.

Parameters:

transforms – The transforms to apply sequentially.

__call__(sample: Sample, *, training: bool = True, generator: torch.Generator | None = None) → Sample

Apply each contained transform in order (each honours its own gate).

apply(sample: Sample, generator: torch.Generator | None) → Sample

Apply every contained transform (train gate handled per-transform).

omniloader.transforms.base.broadcast_mask(mask: torch.Tensor, value: torch.Tensor) → torch.Tensor

Broadcast a validity mask to the value’s shape for torch.where.

Parameters:

• mask – Mask of shape () (vector) or (T,) (sequence).

• value – Value of shape (), (F,), (T,) or (T, F).

Returns:

A boolean tensor broadcastable to value.shape.

omniloader.transforms.base.resolve_keys(keys: collections.abc.Sequence[str] | None, schema: omniloader.schema.spec.UnifiedSchema | None) → list[str]

Return explicit keys or all feature keys from schema.

omniloader.transforms.base.resolve_sequence_keys(keys: collections.abc.Sequence[str] | None, schema: omniloader.schema.spec.UnifiedSchema | None) → list[str]

Return explicit keys or all sequence feature keys from schema.