Avoid use of deprecated device_buffer attriutes of jax.Array

These have been deprecated as of JAX v0.4.22 PiperOrigin-RevId: 589254584
google · Dec 11, 2023 · ad3d524 · ad3d524
1 parent 429dc13
commit ad3d524
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/neural_tangents/_src/predict.py b/neural_tangents/_src/predict.py
@@ -1347,7 +1347,10 @@ def _is_on_cpu(x: PyTree) -> bool:
  def _arr_is_on_cpu(x: jnp.ndarray) -> bool:
  # TODO(romann): revisit when https://github.com/google/jax/issues/1431 and
  # https://github.com/google/jax/issues/1432 are fixed.
- if hasattr(x, 'device_buffer'):
+ if hasattr(x, 'addressable_shards'):
+ # device_buffer is deprecated, so try addressable_shards first.
+ return 'cpu' in str(x.addressable_shards[0].device).lower()
+ elif hasattr(x, 'device_buffer'):
  return 'cpu' in str(x.device_buffer.device()).lower()
 
  if isinstance(x, (np.ndarray, jnp.ndarray)):

diff --git a/tests/predict_test.py b/tests/predict_test.py
@@ -857,9 +857,11 @@ def testPredictOnCPU(self):
  def is_on_cpu(x):
  return jax.tree_util.tree_all(
  jax.tree_map(
- lambda x: 'cpu' in str(x.device_buffer.device()
- ).lower(),
- x))
+ lambda x: 'cpu'
+ in str(x.addressable_shards[0].device).lower(),
+ x,
+ )
+ )
 
  self.assertEqual(on_cpu, is_on_cpu(predict_inf))
  self.assertEqual(on_cpu, is_on_cpu(predict_none))