meshgrid and atleast_*d NumPy 2 updates (#11106)

dask/array/creation.py

@@ -22,7 +22,7 @@
     normalize_chunks,
     stack,
 )
-from dask.array.numpy_compat import AxisError
+from dask.array.numpy_compat import NUMPY_GE_200, AxisError
 from dask.array.ufunc import greater_equal, rint
 from dask.array.utils import meta_from_array
 from dask.array.wrap import empty, full, ones, zeros
@@ -501,9 +501,10 @@ def meshgrid(*xi, sparse=False, indexing="xy", **kwargs):
         grid = broadcast_arrays(*grid)
 
     if indexing == "xy" and len(xi) > 1:
-        grid[0], grid[1] = grid[1], grid[0]
+        grid = (grid[1], grid[0], *grid[2:])
 
-    return grid
+    out_type = tuple if NUMPY_GE_200 else list
+    return out_type(grid)
 
 
 def indices(dimensions, dtype=int, chunks="auto"):

dask/array/routines.py

@@ -83,6 +83,8 @@ def atleast_3d(*arys):
     if len(new_arys) == 1:
         return new_arys[0]
     else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
         return new_arys
 
 
@@ -101,6 +103,8 @@ def atleast_2d(*arys):
     if len(new_arys) == 1:
         return new_arys[0]
     else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
         return new_arys
 
 
@@ -117,6 +121,8 @@ def atleast_1d(*arys):
     if len(new_arys) == 1:
         return new_arys[0]
     else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
         return new_arys
 
 

dask/array/tests/test_creation.py

@@ -381,7 +381,7 @@ def test_meshgrid(shapes, chunks, indexing, sparse):
     r_a = np.meshgrid(*xi_a, indexing=indexing, sparse=sparse)
     r_d = da.meshgrid(*xi_d, indexing=indexing, sparse=sparse)
 
-    assert isinstance(r_d, list)
+    assert type(r_d) is type(r_a)
     assert len(r_a) == len(r_d)
 
     for e_r_a, e_r_d, i in zip(r_a, r_d, do):