mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Explicit doc note about device_put* async

Browse Source
This commit is contained in:
Lucas Beyer 2022-04-04 23:38:51 +02:00 committed by GitHub
parent 6825f654b1
commit f7b749c99c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jax/_src/api.py
View File

@ -2657,6 +2657,8 @@ def device_put(x, device: Optional[xc.Device] = None):
For more details on data placement see the
:ref:`FAQ on data placement <faq-data-placement>`.
This function is always asynchronous, i.e. returns immediately.
Returns:
A copy of ``x`` that resides on ``device``.
"""
@ -2674,6 +2676,8 @@ def device_put_sharded(shards: Sequence[Any], devices: Sequence[xc.Device]):
devices: A sequence of :py:class:`Device` instances representing the devices
to which corresponding shards in ``shards`` will be transferred.
This function is always asynchronous, i.e. returns immediately.
Returns:
A ShardedDeviceArray or (nested) Python container thereof representing the
elements of ``shards`` stacked together, with each shard backed by physical
@ -2743,6 +2747,8 @@ def device_put_replicated(x: Any, devices: Sequence[xc.Device]):
devices: A sequence of :py:class:`Device` instances representing the devices
to which ``x`` will be transferred.
This function is always asynchronous, i.e. returns immediately.
Returns:
A ShardedDeviceArray or (nested) Python container thereof representing the
value of ``x`` broadcasted along a new leading axis of size