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

[JAX] Deprecate .block_host_until_ready() in favor of .block_until_ready().

Browse Source 
JAX kept an older name around (.block_host_until_ready()) in parallel with the new name (.block_until_ready()) to avoid breaking users. Deprecate it so we only have one name.

PiperOrigin-RevId: 433228545
This commit is contained in:
Peter Hawkins 2022-03-08 09:14:14 -08:00 committed by jax authors
parent 809156c1fc
commit 08fbd77d90
3 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
CHANGELOG.md
View File

@ -24,6 +24,9 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
## jaxlib 0.3.1 (Unreleased) ## jaxlib 0.3.1 (Unreleased)
* Changes * Changes
* Deprecations:
* The ``.block_host_until_ready()`` method on JAX arrays has been deprecated.
Use ``.block_until_ready()`` instead.
## jax 0.3.1 (Feb 18, 2022) ## jax 0.3.1 (Feb 18, 2022)
* [GitHub * [GitHub

2
jax/_src/device_array.py
View File

@ -138,7 +138,7 @@ class _DeviceArray(DeviceArray): # type: ignore
Returns the buffer object (`self`). Returns the buffer object (`self`).
""" """
self._check_if_deleted() self._check_if_deleted()
self.device_buffer.block_host_until_ready() # pytype: disable=attribute-error self.device_buffer.block_until_ready()
return self return self
@property @property

2
jax/interpreters/pxla.py
View File

@ -653,7 +653,7 @@ def _sda_check_if_deleted(self):
def _sda_block_until_ready(self): def _sda_block_until_ready(self):
self._check_if_deleted() self._check_if_deleted()
for buf in self.device_buffers: for buf in self.device_buffers:
buf.block_host_until_ready() buf.block_until_ready()
return self return self