gm.data.pad

gemma.gm.data.pad(

element: kauldron.ktyping.pytree.PyTree[Array['sequence']],

max_length: int,

*,

truncate: bool = False,

fill_value: int = 0,

axis: int = -1,

) → kauldron.ktyping.pytree.PyTree[Array['max_length']]

Add zeros to the end of the sequence to reach the max length.

Supports padding multiple arrays at once.

Parameters:

• element – The sequence to pad.

• max_length – The max length of the sequence.

• truncate – Whether to truncate the sequence to the max length. If False, sequences longer than the max_length will raise an error.

• fill_value – The value to fill the padded sequence with.

• axis – The axis in which to pad the sequence (only -1 supported at the moment).

Returns:

The padded sequence of length max_length.