def _split_into_blocks(x: jnp.ndarray, block_len: int, axis: int) -> jnp.ndarray:
    """Split an input array into blocks of a given `block_len` along the given `axis`. If the dimension length
    is not a multiple of `block_len`, it will be padded first with selected `pad_value`.
    """
    # pad tensor to multiple of block_len
    if x.shape[axis] % block_len != 0:
        x = _pad_to_multiple(x, block_len, axis, pad_value=0)
    num_blocks = x.shape[axis] // block_len
    output_shape = x.shape[:axis] + (num_blocks, block_len) + x.shape[(axis + 1) :]
    return x.reshape(output_shape)
Split an input array into blocks of a given `block_len` along the given `axis`. If the dimension length
    is not a multiple of `block_len`, it will be padded first with selected `pad_value`.
    
dpkd/transformers/src/transformers/models/longt5/modeling_flax_longt5.py