Remove xnp.ndarray calls

PiperOrigin-RevId: 515296466

Author: Conchylicultor

dataclass_array/array_dataclass_test.py

@@ -72,8 +72,8 @@ def assert_val(p: Point, shape: Shape, xnp: enp.NpModule = None):
     assert p.y.shape == shape
     assert enp.lazy.as_dtype(p.x.dtype) == np.float32
     assert enp.lazy.as_dtype(p.y.dtype) == np.float32
-    assert isinstance(p.x, xnp.ndarray)
-    assert isinstance(p.y, xnp.ndarray)
+    assert enp.compat.is_array_xnp(p.x, xnp)
+    assert enp.compat.is_array_xnp(p.y, xnp)
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
@@ -100,8 +100,8 @@ def assert_val(p: Isometrie, shape: Shape, xnp: enp.NpModule = None):
     assert p.t.shape == shape + (2,)
     assert enp.lazy.as_dtype(p.r.dtype) == np.float32
     assert enp.lazy.as_dtype(p.t.dtype) == np.int32
-    assert isinstance(p.r, xnp.ndarray)
-    assert isinstance(p.t, xnp.ndarray)
+    assert enp.compat.is_array_xnp(p.r, xnp)
+    assert enp.compat.is_array_xnp(p.t, xnp)
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
@@ -228,8 +228,8 @@ def assert_val(p: WithStatic, shape: Shape, xnp: enp.NpModule = None):
     assert p.y.shape == shape + (2, 2)
     assert enp.lazy.as_dtype(p.x.dtype) == np.float32
     assert enp.lazy.as_dtype(p.y.dtype) == np.float32
-    assert isinstance(p.x, xnp.ndarray)
-    assert isinstance(p.y, xnp.ndarray)
+    assert enp.compat.is_array_xnp(p.x, xnp)
+    assert enp.compat.is_array_xnp(p.y, xnp)
     # Static field is correctly forwarded
     assert isinstance(p.static, str)
     assert p.static == 'abc'
@@ -586,7 +586,7 @@ def test_broadcast(xnp: enp.NpModule):
 def test_infer_np(xnp: enp.NpModule):
   p = Point(x=xnp.ones((3,)), y=[0, 0, 0])  # y is casted to xnp
   assert p.xnp is xnp
-  assert isinstance(p.y, xnp.ndarray)
+  assert enp.compat.is_array_xnp(p.y, xnp)
 
 
 @parametrize_dataclass_arrays

dataclass_array/vectorization.py

@@ -294,7 +294,7 @@ def _broacast_and_flatten_to(
   final_shape = batch_shape + inner_shape
   if isinstance(array, array_dataclass.DataclassArray):
     array = array.broadcast_to(final_shape)
-  elif isinstance(array, xnp.ndarray):
+  elif enp.compat.is_array_xnp(array, xnp):
     array = xnp.broadcast_to(array, final_shape)
   else:
     raise TypeError(f'Unexpected array type: {type(array)}')

dataclass_array/vectorization_test.py

@@ -70,7 +70,7 @@ def test_broadcast_args(
   def fn(self, arg_dc, arg_array):
     assert isinstance(self, dca.testing.Ray)
     assert isinstance(arg_dc, dca.testing.Ray)
-    assert isinstance(arg_array, xnp.ndarray)
+    assert enp.compat.is_array_xnp(arg_array, xnp)
     assert self.shape == ()  # pylint: disable=g-explicit-bool-comparison
     assert arg_dc.shape == expected_arg_shape[1:]
     assert arg_array.shape == expected_arg_shape[1:] + (3,)