jax.Array.to_device

abstract Array.to_device(device, *, stream=None)

Return a copy of the array on the specified device

Parameters:

• device (xc.Device | Sharding) – Device or Sharding to which the created array will be committed.

• stream (int | Any | None) – not implemented, passing a non-None value will lead to an error.

• self (Array)

Returns:

copy of array placed on the specified device or devices.