jax.device_put#

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

将 x 传输到 device。

参数:

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

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

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

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

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

返回:

驻留在 device 上的 x 的副本。

如果 device 参数为 None，则如果操作数已在任何设备上，则此操作的行为类似于恒等函数，否则它会将数据传输到默认设备，未提交。

有关数据放置的更多详细信息，请参阅 关于数据放置的 FAQ。

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