# Path: flaxformer/types.py



from typing import Optional
from flaxformer.types import Array
import jax.numpy as jnp
import numpy as np
    array: Array with rank >= 1 to pad.
    factor: Positive integer factor to pad for.
    axis: A valid axis in `array` to pad.
    mode: The padding mode to use according to `jnp.pad`. Defaults to
      'constant'.  See `jax.numpy.pad` documentation for more.
    constant_values: For 'constant' mode, the pad value to use within `jnp.pad`.
      Defaults to 0.

  Returns:
    The padded Array result.
  """
  array = jnp.asarray(array)

  if factor < 1:
    raise ValueError(f'`factor` must be positive but got {factor}.')
  rank = array.ndim
  if axis < -rank or axis >= rank:
    raise ValueError(
        f'`axis` ({axis}) out of bounds for `array` rank ({rank}).')

  axis_len = array.shape[axis]
  pad_len = -axis_len % factor
  pad_width = [(0, 0)] * rank
  pad_width[axis] = (0, pad_len)
  kwargs = {}
  if mode == 'constant':
    kwargs['constant_values'] = constant_values
  return jnp.pad(array=array, pad_width=pad_width, mode=mode, **kwargs)



# Path: flaxformer/types.py



# Path: flaxformer/architectures/longt5/tensor_utils.py
from typing import Optional
from flaxformer.types import Array
import jax.numpy as jnp
import numpy as np

    array: Array with rank >= 1 to pad.
    factor: Positive integer factor to pad for.
    axis: A valid axis in `array` to pad.
    mode: The padding mode to use according to `jnp.pad`. Defaults to
      'constant'.  See `jax.numpy.pad` documentation for more.
    constant_values: For 'constant' mode, the pad value to use within `jnp.pad`.
      Defaults to 0.

  Returns:
    The padded Array result.
  """
  array = jnp.asarray(array)

  if factor < 1:
    raise ValueError(f'`factor` must be positive but got {factor}.')
  rank = array.ndim
  if axis < -rank or axis >= rank:
    raise ValueError(
        f'`axis` ({axis}) out of bounds for `array` rank ({rank}).')

  axis_len = array.shape[axis]
  pad_len = -axis_len % factor
  pad_width = [(0, 0)] * rank
  pad_width[axis] = (0, pad_len)
  kwargs = {}
  if mode == 'constant':
    kwargs['constant_values'] = constant_values
  return jnp.pad(array=array, pad_width=pad_width, mode=mode, **kwargs)



def split_into_blocks(array: Array,