jax.Array.to_device#

abstractArray.to_device(device,*,stream=None)[source]#

Return a copy of the array on the specified device

Parameters:

• device (xc.Device |Sharding) –Device orShardingto 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.