jax_astype

probly.representation.jax_functions.jax_astype(x: JaxArrayLike, /, dtype: DTypeLike | None, *, copy: bool = False, device: jax.Device | Sharding | None = None) → jax.Array[source]

Cast the array to a new data type, mirroring jax.numpy.astype.

Parameters:

  • x – The array to cast.

  • dtype – The target data type.

  • copy – Whether to always return a copy.

  • device – The device the result should live on.

Returns:

The cast array.