Add jax.Array public type which replaces jnp.DeviceArray and `jax…

….core.Tracer`. PiperOrigin-RevId: 482008250
google · Oct 18, 2022 · 3680144 · 3680144
1 parent 3c4663b
commit 3680144
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/jax_cfd/base/grids.py b/jax_cfd/base/grids.py
@@ -83,8 +83,7 @@ def dtype(self):
   def shape(self) -> Tuple[int, ...]:
     return self.data.shape
 
-  _HANDLED_TYPES = (numbers.Number, np.ndarray, jnp.DeviceArray,
-                    jax.ShapedArray, jax.core.Tracer)
+  _HANDLED_TYPES = (numbers.Number, np.ndarray, jax.Array, jax.ShapedArray)
 
   def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
     """Define arithmetic on GridArrays using NumPy's mixin."""