"""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.
"""
import torch
[docs]
def 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.
Args:
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.
"""
current_size = tensor.shape[0]
if current_size < target_size:
pad_amount = target_size - current_size
pad = tensor.new_full((pad_amount, *tensor.shape[1:]), pad_value)
result = torch.cat([tensor, pad], dim=0)
mask = torch.cat(
[
torch.ones(current_size, dtype=torch.bool),
torch.zeros(pad_amount, dtype=torch.bool),
]
)
elif current_size > target_size:
result = tensor[:target_size]
mask = torch.ones(target_size, dtype=torch.bool)
else:
result = tensor
mask = torch.ones(current_size, dtype=torch.bool)
return result, mask
[docs]
def 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).
Args:
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]])
"""
if tensor.shape[0] == 0:
raise ValueError("Cannot repeat-pad a tensor with zero time steps.")
current = tensor.shape[0]
if current >= target_size:
return tensor
pad_count = target_size - current
last = tensor[-1:].expand(pad_count, *(-1,) * (tensor.ndim - 1))
return torch.cat([tensor, last], dim=0)
[docs]
def 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.
Args:
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
"""
t = tensor.shape[0]
if t == 0:
return tensor
if valid_mask is None:
valid_mask = tensor.reshape(t, -1).any(dim=1)
valid_indices = valid_mask.nonzero(as_tuple=False).squeeze(1)
if valid_indices.numel() == 0:
raise ValueError("Tensor contains no valid time steps — cannot fill gaps.")
result = tensor.clone()
# Back-fill leading zeros from first valid step
first_valid = int(valid_indices[0].item())
if first_valid > 0:
result[:first_valid] = tensor[first_valid]
# Forward-fill remaining gaps
last_valid_idx = first_valid
for i in range(first_valid + 1, t):
if valid_mask[i]:
last_valid_idx = i
else:
result[i] = result[last_valid_idx]
return result