jax.lax.pad

jax.lax.pad(operand, padding_value, padding_config)

Applies low, high, and/or interior padding to an array.

Wraps XLA’s Pad operator.

Parameters:

• operand (jax.typing.ArrayLike) – an array to be padded.

• padding_value (jax.typing.ArrayLike) – the value to be inserted as padding. Must have the same dtype as operand.

• padding_config (Sequence[tuple[int, int, int]]) – a sequence of (low, high, interior) tuples of integers, giving the amount of low, high, and interior (dilation) padding to insert in each dimension.

Returns:

The operand array with padding value padding_value inserted in each dimension according to the padding_config.

Return type:

Array