Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-05-31 04:45:31 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.

Browse Source 
In most cases, use the public jax.* name instead.

PiperOrigin-RevId: 395812511
This commit is contained in:
Peter Hawkins
2021-09-10 00:32:39 +01:00
committed by Saran Tunyasuvunakool
parent 96a13847c6
commit a7d75013c9
2 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
byol/utils/helpers.py
+1 -1
View File
@@ -119,7 +119,7 @@ def bcast_local_devices(value):
def _replicate(x):
"""Replicate an object on each device."""
x = jnp.array(x)
return jax.api.device_put_sharded(len(devices) * [x], devices)
return jax.device_put_sharded(len(devices) * [x], devices)
return jax.tree_util.tree_map(_replicate, value)