Repository: dask/dask

Issue:
arange crashes for parameters near 2**63
python 3.11
numpy 2.2.2
dask 2025.1.0

```python
>>> start, stop, step = 0., -9_131_138_316_486_228_481, -92_233_720_368_547_759
>>> da.arange(start, stop, step)

NotImplementedError: An error occurred while calling the arange method registered to the numpy backend.
Original Message: Can not use auto rechunking with object dtype. We are unable to estimate the size in bytes of object data
```

The problem is that da.arange internally multiplies step by 100, which sends it beyond the limit of np.int64, which in turn triggers edge case behaviour in numpy:

```python
>>> np.arange(start, stop, step * 100)
array([0.0], dtype=object)
```

The last person to ever touch this was... myself, a whopping 7 years ago: https://github.com/dask/dask/pull/3722


Discussion:

Patch:
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 079c5afef..a47ab0c79 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -27,6 +27,7 @@ dependencies:
   - numpy>=2
   - pandas
   - numba
+  - numbagg  # only here to test others with numpy
   - flask
   - h5py
   # Temporarily removing to allow `numpy >=2` to be installed
diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index a67a9c973..37908ca3a 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -25,7 +25,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
         git+https://github.com/dask/zict \
         git+https://github.com/dask/distributed \
         git+https://github.com/zarr-developers/zarr-python
-    mamba uninstall --force numpy pandas scipy numexpr numba sparse scikit-image h5py
+    mamba uninstall --force numpy pandas scipy numexpr numba sparse scikit-image h5py numbagg
     python -m pip install --no-deps --pre --retries 10 \
         -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \
         numpy \
diff --git a/dask/array/creation.py b/dask/array/creation.py
index b906bc405..ba7e96889 100644
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -382,9 +382,7 @@ def linspace(
 
 
 @array_creation_dispatch.register_inplace("numpy")
-def arange(
-    start=None, /, stop=None, step=1, *, chunks="auto", like=None, dtype=None, **kwargs
-):
+def arange(start=None, /, stop=None, step=1, *, chunks="auto", like=None, dtype=None):
     """
     Return evenly spaced values from `start` to `stop` with step size `step`.
 
@@ -431,6 +429,12 @@ def arange(
     elif stop is None:
         start, stop = 0, start
 
+    # Avoid loss of precision calculating blockstart and blockstop below
+    # when start is a very large int (~2**63) and step is a small float
+    if start != 0 and not np.isclose(start + step - start, step, atol=0):
+        r = arange(0, stop - start, step, chunks=chunks, dtype=dtype, like=like)
+        return r + start
+
     num = int(max(np.ceil((stop - start) / step), 0))
 
     meta = meta_from_array(like) if like is not None else None
@@ -440,16 +444,14 @@ def arange(
 
     chunks = normalize_chunks(chunks, (num,), dtype=dtype)
 
-    if kwargs:
-        raise TypeError("Unexpected keyword argument(s): %s" % ",".join(kwargs.keys()))
-
     name = "arange-" + tokenize((start, stop, step, chunks, dtype))
     dsk = {}
     elem_count = 0
 
     for i, bs in enumerate(chunks[0]):
-        blockstart = start + (elem_count * step)
-        blockstop = start + ((elem_count + bs) * step)
+        blockstart = start + elem_count * step
+        blockstop = start + (elem_count + bs) * step
+
         task = Task(
             (name, i),
             partial(chunk.arange, like=meta),
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 27aafc3f8..0458f583b 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -11,6 +11,7 @@ from numbers import Integral, Number
 from operator import mul
 
 import numpy as np
+from packaging.version import Version
 from tlz import accumulate, drop, pluck
 
 import dask.array as da
@@ -37,6 +38,11 @@ from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils import apply, deepmap, derived_from
 
+try:
+    import numbagg
+except ImportError:
+    numbagg = None
+
 if da._array_expr_enabled():
     from dask.array._array_expr import _tree_reduce, reduction
 else:
@@ -1512,16 +1518,28 @@ def nanmedian(a, axis=None, keepdims=False, out=None):
     if builtins.any(a.numblocks[ax] > 1 for ax in axis):
         a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
 
+    if (
+        numbagg is not None
+        and Version(numbagg.__version__).release >= (0, 7, 0)
+        and a.dtype.kind in "uif"
+        and not keepdims
+    ):
+        func = numbagg.nanmedian
+        kwargs = {}
+    else:
+        func = np.nanmedian
+        kwargs = {"keepdims": keepdims}
+
     result = a.map_blocks(
-        np.nanmedian,
+        func,
         axis=axis,
-        keepdims=keepdims,
         drop_axis=axis if not keepdims else None,
         chunks=(
             [1 if ax in axis else c for ax, c in enumerate(a.chunks)]
             if keepdims
             else None
         ),
+        **kwargs,
     )
 
     result = handle_out(out, result)
@@ -1714,18 +1732,30 @@ def nanquantile(
     if builtins.any(a.numblocks[ax] > 1 for ax in axis):
         a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
 
-    if NUMPY_GE_200:
-        kwargs = {"weights": weights}
+    if (
+        numbagg is not None
+        and Version(numbagg.__version__).release >= (0, 8, 0)
+        and a.dtype.kind in "uif"
+        and weights is None
+        and method == "linear"
+        and not keepdims
+    ):
+        func = numbagg.nanquantile
+        kwargs = {"quantiles": q}
     else:
-        kwargs = {}
+        func = _custom_quantile
+        kwargs = {
+            "q": q,
+            "method": method,
+            "interpolation": interpolation,
+            "keepdims": keepdims,
+        }
+        if NUMPY_GE_200:
+            kwargs.update({"weights": weights})
 
     result = a.map_blocks(
-        _custom_quantile,
-        q=q,
-        method=method,
-        interpolation=interpolation,
+        func,
         axis=axis,
-        keepdims=keepdims,
         drop_axis=axis if not keepdims else None,
         new_axis=0 if isinstance(q, Iterable) else None,
         chunks=_get_quantile_chunks(a, q, axis, keepdims),


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index fba157f8a..d0fca5540 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -295,14 +295,17 @@ def test_arange_dtype_force(dtype):
         (0, 2**63 - 1, 2**63 - 10_000),
         (0.0, 2**63 - 1, 2**63 - 10_000),
         (0.0, -9_131_138_316_486_228_481, -92_233_720_368_547_759),
+        (-72_057_594_037_927_945, -72_057_594_037_927_938, 1.0),
+        (-72_057_594_037_927_945, -72_057_594_037_927_938, 1.5),
     ],
 )
-def test_arange_very_large_args(start, stop, step):
+@pytest.mark.parametrize("chunks", ["auto", 1])
+def test_arange_very_large_args(start, stop, step, chunks):
     """Test args that are very close to 2**63
     https://github.com/dask/dask/issues/11706
     """
     a_np = np.arange(start, stop, step)
-    a_da = da.arange(start, stop, step)
+    a_da = da.arange(start, stop, step, chunks=chunks)
     assert_eq(a_np, a_da)