jax.Array.to_device

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

返回指定设备上的数组副本

参数:

• device (xc.Device | Sharding) – Device 或 Sharding 到创建的数组将被提交的位置。

• stream (int | Any | None) – 未实现，传递非None值将导致错误。

• self (Array)

返回:

复制数组到指定的设备或设备上。