jax.numpy.resize

jax.numpy.resize(a, new_shape)

Return a new array with specified shape.

JAX implementation of numpy.resize().

Parameters:

• a (ArrayLike) – input array or scalar.

• new_shape (Shape) – int or tuple of ints. Specifies the shape of the resized array.

Returns:

A resized array with specified shape. The elements of a are repeated in the resized array, if the resized array is larger than the original aray.

Return type:

Array

See also

• jax.numpy.reshape(): Returns a reshaped copy of an array.

• jax.numpy.repeat(): Constructs an array from repeated elements.

Examples

>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> jnp.resize(x, (3, 3))
Array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)
>>> jnp.resize(x, (3, 4))
Array([[1, 2, 3, 4],
       [5, 6, 7, 8],
       [9, 1, 2, 3]], dtype=int32)
>>> jnp.resize(4, (3, 2))
Array([[4, 4],
       [4, 4],
       [4, 4]], dtype=int32, weak_type=True)