jax.numpy.repeat#

jax.numpy.repeat(a, repeats, axis=None, *, total_repeat_length=None)[source]#

Repeat elements of an array.

LAX-backend implementation of numpy.repeat().

JAX adds the optional total_repeat_length parameter which specifies the total number of repeat, and defaults to sum(repeats). It must be specified for repeat to be compilable. If sum(repeats) is larger than the specified total_repeat_length the remaining values will be discarded. In the case of sum(repeats) being smaller than the specified target length, the final value will be repeated.

Original docstring below.

Parameters:

  • a (array_like) – Input array.

  • repeats (int or array of ints) – The number of repetitions for each element. repeats is broadcasted to fit the shape of the given axis.

  • axis (int, optional) – The axis along which to repeat values. By default, use the flattened input array, and return a flat output array.

  • total_repeat_length (Optional[int]) –

Returns:

repeated_array – Output array which has the same shape as a, except along the given axis.

Return type:

ndarray