Return dask array if all axes are squeezed (#9250)

pavithraes · web-flow · commit edfa91e224dd · 2022-07-22T05:53:49.000-07:00
* return 0d array if all axes are squeezed

* update if logic
diff --git a/dask/array/routines.py b/dask/array/routines.py
@@ -1934,6 +1934,13 @@ def squeeze(a, axis=None):
 
     sl = tuple(0 if i in axis else slice(None) for i, s in enumerate(a.shape))
 
+    # Return 0d Dask Array if all axes are squeezed,
+    # to be consistent with NumPy. Ref: https://github.com/dask/dask/issues/9183#issuecomment-1155626619
+    if all(s == 0 for s in sl) and all(s == 1 for s in a.shape):
+        return a.map_blocks(
+            np.squeeze, meta=a._meta, drop_axis=tuple(range(len(a.shape)))
+        )
+
     a = a[sl]
 
     return a
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
@@ -1489,6 +1489,17 @@ def test_squeeze(is_func, axis):
     assert d_s.chunks == exp_d_s_chunks
 
 
+@pytest.mark.parametrize("shape", [(1,), (1, 1)])
+def test_squeeze_1d_array(shape):
+    a = np.full(shape=shape, fill_value=2)
+    a_s = np.squeeze(a)
+    d = da.from_array(a, chunks=(1))
+    d_s = da.squeeze(d)
+    assert isinstance(d_s, da.Array)
+    assert isinstance(d_s.compute(), np.ndarray)
+    assert_eq(d_s, a_s)
+
+
 def test_vstack():
     x = np.arange(5)
     y = np.ones(5)
