Fix the check for jax.Array.

PiperOrigin-RevId: 652787597
Change-Id: I7aff382d61475c35cc24b6bbc42d62b10ccebe76

acme/utils/loggers/base.py

@@ -17,6 +17,7 @@
 import abc
 from typing import Any, Mapping, Optional
 
+import jax
 import numpy as np
 import tree
 from typing_extensions import Protocol
@@ -68,8 +69,8 @@ def close(self):
 def tensor_to_numpy(value: Any):
  if hasattr(value, 'numpy'):
  return value.numpy() # tf.Tensor (TF2).
- if hasattr(value, 'device_buffer'):
- return np.asarray(value) # jnp.DeviceArray.
+ if isinstance(value, jax.Array):
+ return np.asarray(value)
  return value