Repository: dask/dask
Pull request #11327 (2024-08-19 09:00:02)

Fix NumPy overflowing for prod on 2.0
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


phofl: Found here: https://github.com/dask/dask-expr/pull/1125
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ±0       13 suites  ±0   3h 1m 35s :stopwatch: + 1m 32s
 13 158 tests ±0   12 090 :white_check_mark: ± 0   1 068 :zzz: ± 0  0 :x: ±0 
137 107 runs  ±0  117 868 :white_check_mark:  - 37  19 239 :zzz: +37  0 :x: ±0 

Results for commit f9ff4536. ± Comparison against base commit b7d9bf49.

[test-results]:data:application/gzip;base64,H4sIADgRw2YC/02MSw7DIAwFrxKx7sKGEKCXqRKCJdR8KgKrqHcv+ZUsZ/w8KyM/uIU9KxSPii3Jx0J9Cm3087QxaCOzytd43FHqi19LsnaTHAwU+faf/bMpQ2r9kF0ZuRDmcJqQpiOtENTJ/zQqvXcOeaYNF+Zyt/TO97Kdx9HHDIwMUS1FY7kFaRtUqmsdqB4E73WH0jinsSZi3x8ViXKTGgEAAA==


Patch:
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index ff4e2de03..c996831a6 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -741,6 +741,27 @@ def map_overlap(
             # note that keys are relabeled to match values in range(x.ndim)
             depth = {n: depth[ax] for n, ax in enumerate(kept_axes)}
             boundary = {n: boundary[ax] for n, ax in enumerate(kept_axes)}
+
+        # add any new axes to depth and boundary variables
+        new_axis = kwargs.pop("new_axis", None)
+        if new_axis is not None:
+            if isinstance(new_axis, Number):
+                new_axis = [new_axis]
+
+            # convert negative new_axis to equivalent positive value
+            ndim_out = max(a.ndim for a in args if isinstance(a, Array))
+            new_axis = [d % ndim_out for d in new_axis]
+
+            for axis in new_axis:
+                for existing_axis in list(depth.keys()):
+                    if existing_axis >= axis:
+                        # Shuffle existing axis forward to give room to insert new_axis
+                        depth[existing_axis + 1] = depth[existing_axis]
+                        boundary[existing_axis + 1] = boundary[existing_axis]
+
+                depth[axis] = 0
+                boundary[axis] = "none"
+
         return trim_internal(x, depth, boundary)
     else:
         return x
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 4780f71d7..a56e9179b 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -345,6 +345,7 @@ def tensordot(lhs, rhs, axes=2):
     if concatenate:
         return intermediate
     else:
+        left_axes = [ax if ax >= 0 else lhs.ndim + ax for ax in left_axes]
         return intermediate.sum(axis=left_axes)
 
 


Test patch:
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index 659085923..42a5c87e8 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -826,3 +826,17 @@ def test_sliding_window_errors(window_shape, axis):
     arr = da.zeros((4, 3))
     with pytest.raises(ValueError):
         sliding_window_view(arr, window_shape, axis)
+
+
+def test_map_overlap_new_axis():
+    arr = da.arange(6, chunks=2)
+    assert arr.shape == (6,)
+    assert arr.chunks == ((2, 2, 2),)
+
+    actual = arr.map_overlap(lambda x: np.stack([x, x + 0.5]), depth=1, new_axis=[0])
+    expected = np.stack([np.arange(6), np.arange(6) + 0.5])
+
+    assert actual.chunks == ((1,), (2, 2, 2))
+    # Shape and chunks aren't known until array is computed,
+    # so don't expclitly check shape or chunks in assert_eq
+    assert_eq(expected, actual, check_shape=False, check_chunks=False)
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 4b1ca1e47..0b07ccc14 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -334,7 +334,7 @@ def test_tensordot():
     y = np.arange(200).reshape((20, 10))
     b = da.from_array(y, chunks=(4, 5))
 
-    for axes in [1, (1, 0)]:
+    for axes in [1, (1, 0), (-1, 0)]:
         assert_eq(da.tensordot(a, b, axes=axes), np.tensordot(x, y, axes=axes))
         assert_eq(da.tensordot(x, b, axes=axes), np.tensordot(x, y, axes=axes))
         assert_eq(da.tensordot(a, y, axes=axes), np.tensordot(x, y, axes=axes))
diff --git a/dask/dataframe/tests/test_ufunc.py b/dask/dataframe/tests/test_ufunc.py
index befb5c45b..e5366a0c8 100644
--- a/dask/dataframe/tests/test_ufunc.py
+++ b/dask/dataframe/tests/test_ufunc.py
@@ -4,6 +4,8 @@ import warnings
 
 import pytest
 
+from dask.array.numpy_compat import NUMPY_GE_200
+
 pd = pytest.importorskip("pandas")
 
 import numpy as np
@@ -519,6 +521,14 @@ def test_ufunc_with_reduction(redfunc, ufunc, pandas):
     np_redfunc = getattr(np, redfunc)
     np_ufunc = getattr(np, ufunc)
 
+    if (
+        NUMPY_GE_200
+        and redfunc == "prod"
+        and ufunc in ("floor", "ceil", "trunc")
+        and isinstance(pandas, pd.DataFrame)
+    ):
+        pytest.skip("Numpy started overflowing while we are casting to float")
+
     if (
         redfunc == "prod"
         and ufunc in ["conj", "square", "negative", "absolute"]