jax.device_put#

jax.device_put(x, device=None, *, src=None, donate=False, may_alias=None)[源代码]#

将 x 传输到 device。

参数:

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

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

  • src (None | xc.Device | Sharding | Layout | Any | TransferToMemoryKind | None) – (可选的) Device、Sharding 或标准 Python 容器中的（嵌套）Sharding (必须是 x 的树前缀)，表示 x 所属的设备。

  • donate (bool | Any) – bool 或标准 Python 容器中的（嵌套）bool (必须是 x 的树前缀)。如果为 True，则可以在调用者中覆盖 x 并将其标记为已删除。这是尽力而为。如果可能，JAX 将捐赠，否则不会。如果捐赠，输入缓冲区（在未来）将始终被删除。

  • may_alias (bool | None | Any | None) – bool 或 None 或标准 Python 容器中的（嵌套）bool (必须是 x 的树前缀)。如果为 False，则将复制 x。如果为 true，则 x 可能会根据运行时的实现进行别名。

返回:

一个驻留在 device 上的 x 的副本。

如果 device 参数为 None，则如果操作数已经位于任何设备上，则此操作的行为类似于标识函数，否则它会将数据传输到默认设备，但不会提交。

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

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