jax.Array.to_device

抽象 Array.to_device(device, *, stream=None)

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

参数:

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

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

• self (Array)

返回:

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