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

Use jax.api.device_put_sharded() in place of private JAX APIs.

Browse Source 
PiperOrigin-RevId: 332514384
This commit is contained in:
Jake VanderPlas
2020-09-18 21:21:27 +01:00
committed by Saran Tunyasuvunakool
parent 1d763e0beb
commit 0e5237df2a
1 changed files with 1 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
byol/utils/helpers.py
+1 -3
View File
@@ -119,9 +119,7 @@ def bcast_local_devices(value):
def _replicate(x): def _replicate(x):
"""Replicate an object on each device.""" """Replicate an object on each device."""
x = jnp.array(x) x = jnp.array(x)
aval = jax.ShapedArray((len(devices),) + x.shape, x.dtype) return jax.api.device_put_sharded(len(devices) * [x], devices)
buffers = [jax.interpreters.xla.device_put(x, d) for d in devices]
return jax.pxla.ShardedDeviceArray(aval, buffers)
return jax.tree_util.tree_map(_replicate, value) return jax.tree_util.tree_map(_replicate, value)