jax_reshape

probly.representation.jax_functions.jax_reshape(a: JaxArrayLike, /, shape: int | Sequence[int], order: str = 'C', *, copy: bool | None = None) → jax.Array[source]

Return a reshaped version of the array, mirroring jax.numpy.reshape.

Parameters:

  • a – The array to reshape.

  • shape – The target shape.

  • order – The read/write element order.

  • copy – Whether to copy the data.

Returns:

The reshaped array.