omniloader.utils.padding

Padding and cropping utilities for the time dimension.

Adapted from exordium.utils.padding so OmniLoader can bring variable-length sequences to a fixed target length and produce the matching validity masks.

Functions

pad_or_crop_time_dim(→ tuple[torch.Tensor, torch.Tensor])	Adjust the time dimension (T) of a tensor to a given size and generate a mask.
repeat_pad_time_dim(→ torch.Tensor)	Pad the time dimension by repeating the last time step.
fill_gaps_with_repeat(→ torch.Tensor)	Fill zero-valued time steps by repeating the nearest earlier valid step.

Module Contents

omniloader.utils.padding.pad_or_crop_time_dim(tensor: torch.Tensor, target_size: int, pad_value: int | float = 0) → tuple[torch.Tensor, torch.Tensor]

Adjust the time dimension (T) of a tensor to a given size and generate a mask.

Operates on axis 0 only and is agnostic to the rank of the trailing axes, so it handles (T,), (T, F) and structured sequences such as (T, C, H, W) alike.

Parameters:

• tensor – Input tensor with shape (T, *trailing).

• target_size – Desired size for the time dimension.

• pad_value – Value to use for padding. Defaults to 0.

Returns:

Tuple of (adjusted_tensor, mask) where mask is a bool tensor of shape (target_size,) with True for original values and False for padded positions.

omniloader.utils.padding.repeat_pad_time_dim(tensor: torch.Tensor, target_size: int) → torch.Tensor

Pad the time dimension by repeating the last time step.

If target_size <= T the tensor is returned unchanged (no cropping).

Parameters:

• tensor – Input tensor of shape (T, ...) where T is the time dimension (e.g. (T, C, H, W) or (T, D)).

• target_size – Desired length of the time dimension.

Returns:

Tensor of shape (target_size, ...) where steps beyond the original length are copies of the last time step.

Raises:

ValueError – If the tensor has zero time steps.

Example:

>>> x = torch.tensor([[1, 2], [3, 4]])  # (2, 2)
>>> repeat_pad_time_dim(x, 5)
tensor([[1, 2], [3, 4], [3, 4], [3, 4], [3, 4]])

omniloader.utils.padding.fill_gaps_with_repeat(tensor: torch.Tensor, valid_mask: torch.Tensor | None = None) → torch.Tensor

Fill zero-valued time steps by repeating the nearest earlier valid step.

Given a tensor of shape (T, ...) where some time steps are entirely zero (no data), this function fills each gap by copying the closest preceding valid step. If the tensor starts with zeros before the first valid step, those leading positions are back-filled from the first valid step.

A time step is considered valid when any element is non-zero, unless an explicit valid_mask is provided.

Parameters:

• tensor – Input tensor of shape (T, ...).

• valid_mask – Optional boolean tensor of shape (T,). True marks valid (non-gap) time steps. When None, validity is inferred as tensor.flatten(1).any(dim=1).

Returns:

Tensor of the same shape with all gaps filled.

Raises:

ValueError – If the tensor contains no valid time steps at all.

Example:

>>> x = torch.zeros(6, 2)
>>> x[1] = torch.tensor([1.0, 2.0])
>>> x[4] = torch.tensor([3.0, 4.0])
>>> fill_gaps_with_repeat(x)
tensor([[1., 2.],   # back-filled from step 1
        [1., 2.],   # valid
        [1., 2.],   # forward-filled from step 1
        [1., 2.],   # forward-filled from step 1
        [3., 4.],   # valid
        [3., 4.]])  # forward-filled from step 4