Repository: dask/dask
Pull request #10636 (2023-11-18 00:13:03)

Update gpuCI `RAPIDS_VER` to `24.02`
New cuDF and ucx-py nightly versions have been detected.

Updated `axis.yaml` to use `24.02`.


Patch:
diff --git a/continuous_integration/gpuci/axis.yaml b/continuous_integration/gpuci/axis.yaml
index dd7525507..c66f47166 100644
--- a/continuous_integration/gpuci/axis.yaml
+++ b/continuous_integration/gpuci/axis.yaml
@@ -9,6 +9,6 @@ LINUX_VER:
 - ubuntu20.04
 
 RAPIDS_VER:
-- "23.12"
+- "24.02"
 
 excludes:
diff --git a/dask/array/linalg.py b/dask/array/linalg.py
index d363a1982..d0112671f 100644
--- a/dask/array/linalg.py
+++ b/dask/array/linalg.py
@@ -1482,11 +1482,10 @@ def norm(x, ord=None, axis=None, keepdims=False):
         if len(axis) == 1:
             raise ValueError("Invalid norm order for vectors.")
 
-    # Coerce to double precision.
-    r = x.astype(np.promote_types(x.dtype, float))
+    r = abs(x)
 
     if ord is None:
-        r = (abs(r) ** 2).sum(axis=axis, keepdims=keepdims) ** 0.5
+        r = (r**2).sum(axis=axis, keepdims=keepdims) ** 0.5
     elif ord == "nuc":
         if len(axis) == 1:
             raise ValueError("Invalid norm order for vectors.")
@@ -1495,7 +1494,6 @@ def norm(x, ord=None, axis=None, keepdims=False):
 
         r = svd(x)[1][None].sum(keepdims=keepdims)
     elif ord == np.inf:
-        r = abs(r)
         if len(axis) == 1:
             r = r.max(axis=axis, keepdims=keepdims)
         else:
@@ -1503,7 +1501,6 @@ def norm(x, ord=None, axis=None, keepdims=False):
             if keepdims is False:
                 r = r.squeeze(axis=axis)
     elif ord == -np.inf:
-        r = abs(r)
         if len(axis) == 1:
             r = r.min(axis=axis, keepdims=keepdims)
         else:
@@ -1516,7 +1513,6 @@ def norm(x, ord=None, axis=None, keepdims=False):
 
         r = (r != 0).astype(r.dtype).sum(axis=axis, keepdims=keepdims)
     elif ord == 1:
-        r = abs(r)
         if len(axis) == 1:
             r = r.sum(axis=axis, keepdims=keepdims)
         else:
@@ -1524,7 +1520,7 @@ def norm(x, ord=None, axis=None, keepdims=False):
             if keepdims is False:
                 r = r.squeeze(axis=axis)
     elif len(axis) == 2 and ord == -1:
-        r = abs(r).sum(axis=axis[0], keepdims=True).min(axis=axis[1], keepdims=True)
+        r = r.sum(axis=axis[0], keepdims=True).min(axis=axis[1], keepdims=True)
         if keepdims is False:
             r = r.squeeze(axis=axis)
     elif len(axis) == 2 and ord == 2:
@@ -1539,6 +1535,6 @@ def norm(x, ord=None, axis=None, keepdims=False):
         if len(axis) == 2:
             raise ValueError("Invalid norm order for matrices.")
 
-        r = (abs(r) ** ord).sum(axis=axis, keepdims=keepdims) ** (1.0 / ord)
+        r = (r**ord).sum(axis=axis, keepdims=keepdims) ** (1.0 / ord)
 
     return r
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 360cc592d..1408627c3 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -4380,12 +4380,19 @@ Dask Name: {name}, {layers}""".format(
         )
 
     @derived_from(pd.Series)
-    def clip(self, lower=None, upper=None, out=None):
+    def clip(self, lower=None, upper=None, out=None, axis=None):
         if out is not None:
             raise ValueError("'out' must be None")
+        if axis not in (None, 0):
+            raise ValueError(f"Series.clip does not support axis={axis}")
         # np.clip may pass out
         return self.map_partitions(
-            M.clip, lower=lower, upper=upper, enforce_metadata=False
+            M.clip,
+            lower=lower,
+            upper=upper,
+            enforce_metadata=False,
+            # See: https://github.com/rapidsai/cudf/issues/14369
+            **({} if axis is None else {"axis": axis}),
         )
 
     @derived_from(pd.Series)
@@ -5662,11 +5669,16 @@ class DataFrame(_Frame):
         return self.map_partitions(M.dropna, **kwargs, enforce_metadata=False)
 
     @derived_from(pd.DataFrame)
-    def clip(self, lower=None, upper=None, out=None):
+    def clip(self, lower=None, upper=None, out=None, axis=None):
         if out is not None:
             raise ValueError("'out' must be None")
         return self.map_partitions(
-            M.clip, lower=lower, upper=upper, enforce_metadata=False
+            M.clip,
+            lower=lower,
+            upper=upper,
+            enforce_metadata=False,
+            # See: https://github.com/rapidsai/cudf/issues/14369
+            **({} if axis is None else {"axis": axis}),
         )
 
     @derived_from(pd.DataFrame)
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index c9e276694..438fa81cb 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -17,6 +17,35 @@ Zero-copy P2P Array Rechunking
 Users should see significant performance improvements when using in-memory P2P array rechunking.
 This is due to no longer copying underlying data buffers.
 
+Below shows a simple example where we compare performance of different rechunking methods.
+
+.. code:: python
+
+  shape = (30_000, 6_000, 150) # 201.17 GiB
+  input_chunks = (60, -1, -1) # 411.99 MiB
+  output_chunks = (-1, 6, -1) # 205.99 MiB
+
+  arr = da.random.random(size, chunks=input_chunks)
+  with dask.config.set({
+      "array.rechunk.method": "p2p",
+      "distributed.p2p.disk": True,
+  }):
+      (
+        da.random.random(size, chunks=input_chunks)
+        .rechunk(output_chunks)
+        .sum()
+        .compute()
+      )
+
+.. image:: images/changelog/2023110-rechunking-disk-perf.png
+  :width: 75%
+  :align: center
+  :alt: A comparison of rechunking performance between the different methods
+    tasks, p2p with disk and p2p without disk on different cluster sizes. The
+    graph shows that p2p without disk is up to 60% faster than the default
+    tasks based approach.
+
+
 See :pr-distributed:`8282`, :pr-distributed:`8318`, :pr-distributed:`8321` from `crusaderky`_ and
 (:pr-distributed:`8322`) from `Hendrik Makait`_ for details.
 
@@ -211,7 +240,7 @@ the ``override_with=`` keyword (see :issue:`10519`).
 This release restores the previous behavior.
 
 See :pr:`10521` from `crusaderky`_ for details.
-    
+
 Complex dtypes in Dask Array reductions
 """""""""""""""""""""""""""""""""""""""
 This release includes improved support for using common reductions


Test patch:
diff --git a/dask/array/tests/test_linalg.py b/dask/array/tests/test_linalg.py
index b0c1c512a..d5006571d 100644
--- a/dask/array/tests/test_linalg.py
+++ b/dask/array/tests/test_linalg.py
@@ -1011,6 +1011,32 @@ def test_norm_any_ndim(shape, chunks, axis, norm, keepdims):
     assert_eq(a_r, d_r)
 
 
+@pytest.mark.xfail(
+    _np_version < parse_version("1.23"),
+    reason="https://github.com/numpy/numpy/pull/17709",
+    strict=False,
+)
+@pytest.mark.parametrize("precision", ["single", "double"])
+@pytest.mark.parametrize("isreal", [True, False])
+@pytest.mark.parametrize("keepdims", [False, True])
+@pytest.mark.parametrize("norm", [None, 1, -1, np.inf, -np.inf])
+def test_norm_any_prec(norm, keepdims, precision, isreal):
+    shape, chunks, axis = (5,), (2,), None
+
+    precs_r = {"single": "float32", "double": "float64"}
+    precs_c = {"single": "complex64", "double": "complex128"}
+
+    dtype = precs_r[precision] if isreal else precs_c[precision]
+
+    a = np.random.default_rng().random(shape).astype(dtype)
+    d = da.from_array(a, chunks=chunks)
+    d_a = np.linalg.norm(a, ord=norm, axis=axis, keepdims=keepdims)
+    d_r = da.linalg.norm(d, ord=norm, axis=axis, keepdims=keepdims)
+
+    assert d_r.dtype == precs_r[precision]
+    assert d_r.dtype == d_a.dtype
+
+
 @pytest.mark.slow
 @pytest.mark.xfail(
     sys.platform == "darwin" and _np_version < parse_version("1.22"),
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index bef2676e5..a459e11e8 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -849,6 +849,50 @@ def test_clip(lower, upper):
     assert_eq(ds.clip(upper=upper), s.clip(upper=upper))
 
 
+def test_clip_axis_0():
+    df = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+            "b": [3, 5, 2, 5, 7, 2, 4, 2, 4],
+        }
+    )
+    s = pd.Series([1, 2, 3, 4, 5, 6, 7, 8, 9])
+    l = pd.Series([3] * len(df))
+    u = pd.Series([7] * len(df))
+
+    ddf = dd.from_pandas(df, 3)
+    ds = dd.from_pandas(s, 3)
+    dl = dd.from_pandas(l, 3)
+    du = dd.from_pandas(u, 3)
+
+    assert_eq(ddf.clip(lower=dl, upper=du, axis=0), df.clip(lower=l, upper=u, axis=0))
+    assert_eq(ddf.clip(lower=dl, axis=0), df.clip(lower=l, axis=0))
+    assert_eq(ddf.clip(upper=du, axis=0), df.clip(upper=u, axis=0))
+
+    assert_eq(ds.clip(lower=dl, upper=du, axis=0), s.clip(lower=l, upper=u, axis=0))
+    assert_eq(ds.clip(lower=dl, axis=0), s.clip(lower=l, axis=0))
+    assert_eq(ds.clip(upper=du, axis=0), s.clip(upper=u, axis=0))
+
+
+def test_clip_axis_1():
+    df = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+            "b": [3, 5, 2, 5, 7, 2, 4, 2, 4],
+        }
+    )
+    ddf = dd.from_pandas(df, 3)
+    l = pd.Series({"a": 2, "b": 3})
+    u = pd.Series({"a": 7, "b": 5})
+
+    assert_eq(ddf.clip(lower=l, upper=u, axis=1), df.clip(lower=l, upper=u, axis=1))
+    assert_eq(ddf.clip(lower=l, axis=1), df.clip(lower=l, axis=1))
+    assert_eq(ddf.clip(upper=u, axis=1), df.clip(upper=u, axis=1))
+
+    with pytest.raises(ValueError, match="Series.clip does not support axis=1"):
+        ddf.a.clip(lower=l, upper=u, axis=1)
+
+
 def test_squeeze():
     df = pd.DataFrame({"x": [1, 3, 6]})
     df2 = pd.DataFrame({"x": [0]})