Replace array by asarray for better Pytorch compatibility

PiperOrigin-RevId: 515265898

Author: Conchylicultor

dataclass_array/array_dataclass_test.py

@@ -280,8 +280,8 @@ def test_point_infered_np(
     shape: Shape,
 ):
   if xnp is not None:  # Normalize np arrays to test the various backend
-    x = xnp.array(x)
-    y = xnp.array(y)
+    x = xnp.asarray(x)
+    y = xnp.asarray(y)
   else:
     xnp = np
 
@@ -565,7 +565,7 @@ def test_broadcast(xnp: enp.NpModule):
   p = Nested(
       # pt.shape broadcasted to (2, 3)
       pt=Point(
-          x=xnp.array(0),
+          x=xnp.asarray(0),
           y=xnp.zeros(broadcast_shape + (3,)),
       ),
       # iso.shape == (), broadcasted to (2, 3)
@@ -642,13 +642,13 @@ class PointNoCast(dca.DataclassArray):
 
   with pytest.raises(ValueError, match='Cannot cast float16'):
     PointNoCast(
-        x=xnp.array([1, 2, 3], dtype=np.float16),
-        y=xnp.array([1, 2, 3], dtype=np.float16),
+        x=xnp.asarray([1, 2, 3], dtype=np.float16),
+        y=xnp.asarray([1, 2, 3], dtype=np.float16),
     )
 
   p = PointNoCast(
-      x=xnp.array([1, 2, 3], dtype=np.float16),
-      y=xnp.array([1, 2, 3], dtype=np.uint8),
+      x=xnp.asarray([1, 2, 3], dtype=np.float16),
+      y=xnp.asarray([1, 2, 3], dtype=np.uint8),
   )
   assert p.shape == (3,)
   assert enp.lazy.as_dtype(p.x.dtype) == np.float16
@@ -665,7 +665,7 @@ class PointNoList(dca.DataclassArray):
 
   with pytest.raises(TypeError, match='Could not infer numpy module'):
     PointNoList(
-        x=xnp.array(1, dtype=np.float16),
+        x=xnp.asarray(1, dtype=np.float16),
         y=[1, 2, 3],
     )
 
@@ -679,8 +679,8 @@ class PointNoBroadcast(dca.DataclassArray):
 
   with pytest.raises(ValueError, match='Cannot broadcast'):
     PointNoBroadcast(
-        x=xnp.array(1, dtype=np.float16),
-        y=xnp.array([1, 2, 3], dtype=np.int32),
+        x=xnp.asarray(1, dtype=np.float16),
+        y=xnp.asarray([1, 2, 3], dtype=np.int32),
     )
 
 

dataclass_array/utils/np_utils_test.py

@@ -29,10 +29,12 @@
 @enp.testing.parametrize_xnp()
 def test_get_xnp(xnp: enp.NpModule):
   # Dataclass array support
-  r = dca.testing.Ray(pos=xnp.array([3.0, 0, 0]), dir=xnp.array([3.0, 0, 0]))
+  r = dca.testing.Ray(
+      pos=xnp.asarray([3.0, 0, 0]), dir=xnp.asarray([3.0, 0, 0])
+  )
   assert np_utils.get_xnp(r) is xnp
   # Array support
-  assert np_utils.get_xnp(xnp.array([3.0, 0, 0])) is xnp
+  assert np_utils.get_xnp(xnp.asarray([3.0, 0, 0])) is xnp
 
   with pytest.raises(TypeError, match='Unexpected array type'):
     np_utils.get_xnp('abc')