jax.Array.to_device

jax.Array.to_device#

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

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

参数:
  • device (xc.Device | Sharding) – 创建的数组将被提交到的 DeviceSharding.

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

  • self (Array)

返回:

放置在指定设备或设备上的数组副本。