jax.设备放置

jax.device_put(x, device=None, *, src=None)

将 x 转移到 device。

参数:

• x – 一个数组、标量或（嵌套的）标准 Python 容器。

• device (None | xc.Device | Sharding | Layout | Any | TransferToMemoryKind) – 可选的 设备、分片，或标准 Python 容器中的（嵌套）:py:class:分片`（必须是 ``x` 的树前缀），表示 x 应传输到的设备。如果给出，则结果将提交到设备。

• src (None | xc.Device | Sharding | Layout | Any | TransferToMemoryKind)

返回:

位于 device 上的 x 的副本。

如果 device 参数是 None，那么如果操作数已经在任何设备上，此操作的行为就像恒等函数一样，否则它会将数据传输到默认设备，未提交。

有关数据放置的更多详细信息，请参阅 数据放置常见问题解答。

此函数始终是异步的，即立即返回，不会阻塞调用它的Python线程，直到任何传输完成。