def _pad_to_multiple(x: jnp.ndarray, block_len: int, axis: int, pad_value: int = 0) -> jnp.ndarray:
    """Pad an array so that a sequence length will be a multiple of `block_len`"""
    pad_len = -x.shape[axis] % block_len
    pad = [(0, 0)] * x.ndim
    pad[axis] = (0, pad_len)
    x = jnp.pad(x, pad_width=pad, mode="constant", constant_values=pad_value)
    return x
Pad an array so that a sequence length will be a multiple of `block_len`
dpkd/transformers/src/transformers/models/longt5/modeling_flax_longt5.py