Batch and Pad Variable-Length Vectors with Optional Final-Batch Dropping

Problem Statement

In machine learning pipelines processing sequential telemetry, physical sensor logs, or embeddings, inputs often arrive as ordered sequences of variable lengths. Prior to feeding these inputs into tensor-based models, they must be organized into batches and padded.

Write a function batch_and_pad that takes a list of variable-length floating-point vectors and produces a list of padded 2D batches according to the following rules:

1. Batching: Partition the ordered input sequences into contiguous batches of size batch_size.
2. Incomplete Final Batch: If drop_last is True and the final batch contains strictly fewer than batch_size elements, discard that final batch. If drop_last is False, retain it.
3. Batch-Local Right Padding: For each batch, determine the maximum length L_max among all sequences within that specific batch. Pad each sequence on the right with pad_value until its length equals L_max.
4. Empty Sequence Handling: If an input sequence is empty ([]), it should still be treated as a valid vector of length 0. If an entire batch contains only empty sequences, L_max is 0, and each sequence in that batch remains empty.

Function Signature

def batch_and_pad(
    sequences: list[list[float]],
    batch_size: int,
    pad_value: float = 0.0,
    drop_last: bool = False
) -> list[list[list[float]]]:
    pass

Examples

Example 1

sequences = [[1.0, 2.0], [3.0], [], [4.0, 5.0, 6.0], [7.0]]
batch_size = 2
pad_value = 0.0
drop_last = False

output = batch_and_pad(sequences, batch_size, pad_value, drop_last)
# Output:
# [
#   [[1.0, 2.0], [3.0, 0.0]],
#   [[0.0, 0.0, 0.0], [4.0, 5.0, 6.0]],
#   [[7.0]]
# ]

Explanation:

• Batch 1 has vectors of lengths 2 and 1 (L_max = 2). [3.0] is right-padded with one 0.0.
• Batch 2 has vectors of lengths 0 and 3 (L_max = 3). [] is right-padded with three 0.0s.
• Batch 3 contains the remaining element [7.0]. Since drop_last=False, it is kept.

Example 2

sequences = [[1.0, 2.0], [3.0], [], [4.0, 5.0, 6.0], [7.0]]
batch_size = 2
pad_value = -1.0
drop_last = True

output = batch_and_pad(sequences, batch_size, pad_value, drop_last)
# Output:
# [
#   [[1.0, 2.0], [3.0, -1.0]],
#   [[-1.0, -1.0, -1.0], [4.0, 5.0, 6.0]]
# ]

Explanation: The final batch containing only [7.0] has size 1 (< batch_size), so it is dropped.

Constraints

• 0 <= len(sequences) <= 10^5
• 0 <= len(seq) <= 5 * 10^3 for each sequence seq in sequences
• Total number of elements across all sequences <= 10^6
• 1 <= batch_size <= 10^4
• pad_value is a finite float.