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

Reverts 6dd1417d4a0a9ee31d8a014352b3a0fb2bcfcbaf

Browse Source 
PiperOrigin-RevId: 691417832
This commit is contained in:
Thomas Köppe 2024-10-30 07:53:24 -07:00 committed by jax authors
parent 2652ab5608
commit 2bed1e88e4
3 changed files with 0 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/_src/api.py
View File

@ -2440,13 +2440,6 @@ def device_put_replicated(x: Any, devices: Sequence[xc.Device]): # noqa: F811
def _device_get(x):
if isinstance(x, core.Tracer):
return x
if dtypes.issubdtype(getattr(x, "dtype", None), dtypes.extended):
try:
to_device = x.dtype._rules.device_get
except AttributeError:
pass
else:
return to_device(x)
try:
toarray = x.__array__
except AttributeError:

5
jax/_src/prng.py
View File

@ -400,11 +400,6 @@ class KeyTyRules:
phys_result = phys_handler(phys_arrays)
return PRNGKeyArray(aval.dtype._impl, phys_result)
@staticmethod
def device_get(val):
buffer = api.device_get(random_unwrap(val))
return random_wrap(buffer, impl=val.dtype._impl)
@staticmethod
def device_put_sharded(vals, aval, sharding, devices):
physical_aval = core.physical_aval(aval)

5
tests/random_test.py
View File

@ -936,11 +936,6 @@ class KeyArrayTest(jtu.JaxTestCase):
x = jnp.array([True, False, False])
f(x) # doesn't crash
def test_device_get(self):
keys = self.make_keys(4)
keys_on_host = jax.device_get(keys)
self.assertKeysEqual(keys, keys_on_host)
def test_device_put(self):
device = jax.devices()[0]
keys = self.make_keys(4)