jax.Array.to_device

jax.Array.to_device#

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

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.

Contents