Bincount fix slicing

[GenevieveBuckley]

We've found problems trying to slice the output of dask.array.bincount() following #7183

• Passes black dask / flake8 dask

• Tests added / passed

  • Added an assertion to the test in dask/array/test/test_routines.py to check the output of da.bincount can be sliced
  • Passes previously failing tests in skimage: pytest skimage/filters/tests/test_thresholding.py::test_thresholds_dask_compatibility -vv
  • Passes previously failing cupy tests
  • Passes previously failing tests in dask-ml

Relevant issues:

• Closes test_thresholding fails with last dask version (2021.3.0) scikit-image/scikit-image#5266
• Closes IndexError("Too many indices for array") raised when attempting to run the K-Means|| example dask-ml#803
• Partially addresses Failing CuPy tests #7324 (some other cupy test failures are unrelated to bincount)
• Closes Compliance with Dask version 2021.03 ? xgcm/xhistogram#27

Summary:
When dask.array.bincount was modified in #7183 we lost the ability to slice the output returned by this function. The output shape and chunks attributes were empty tuples instead of what we expected.

We want to keep the benefits of PR 7183 because it's a lot better at managing memory when given large datasets.

So to add back in the ability to slice the output from bincount, I've added some of the extra logic normally handled by dask.array.reductions.reduction() to keep track of the array shape & chunk sizes.

[GenevieveBuckley]

The CI is currently failing because in my attempt to fix the failing cupy test, I've used the numpy array creation like= kwarg, which was only introduced in numpy version 1.20 (whereas the dask minimum dependency is currently numpy 1.15.1. Discussion here #7324 (comment)

[pentschev]

Thanks @GenevieveBuckley , I can confirm the test passes on my end as well. The fact that we now need NumPy>=1.20 is technically a regression, but probably a small one, I'm fine with telling users they need to upgrade. But could you add the line below to test_bincount to ensure it doesn't fail on older NumPy versions?

@pytest.mark.skipif(not _numpy_120, reason="NEP-35 is not available")

[rabernat]

Thanks @GenevieveBuckley for finding this! In xhistogram (xgcm/xhistogram#27) we are experiencing problems with bincount with dask 2021.03 too.

Specifically:

import dask.array as dsa
ones = dsa.ones(shape=(100,), chunks=(110), dtype='int')
print(dsa.bincount(ones, minlength=2))
# -> dask.array<_bincount_agg-aggregate, shape=(), dtype=int64, chunksize=(), chunktype=numpy.ndarray>

the output of bincount no longer has a shape, so it can't be reshaped (which is what we need to do).

It looks like this PR will fix that. So 🙌 for your work!

[GenevieveBuckley]

It's nice to hear that @rabernat - this has been an unexpectedly high-impact change. I only knew about the problems in scikit-image when I started looking at it.

[GenevieveBuckley]

The cupy bincount test passes locally now (I have cupy version 9.0.0b3 and numpy version 1.20.1

pytest dask/array/tests/test_cupy.py::test_bincount -v

[GenevieveBuckley]

I think this is ready to merge @dask/maintenance

[pentschev]

@GenevieveBuckley thanks for the latest changes, I added another suggestion on making this a bit safer, not sure if that's really possible to happen in Dask but I guess there's no harm in doing so.

[jrbourbeau]

Thanks for your work on this @GenevieveBuckley!

Since axis is always (0,) we can streamline some of the logic around constructing chunks tuples (see the diff below). My hope is that will make reasoning about the logic here easier in the future.

Diff:

diff --git a/dask/array/routines.py b/dask/array/routines.py
index 8d53541e..69f56b63 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -645,7 +645,6 @@ def bincount(x, weights=None, minlength=0, split_every=None):
         if weights.chunks != x.chunks:
             raise ValueError("Chunks of input array x and weights must match.")

-    axis = (0,)
     token = tokenize(x, weights, minlength)
     args = [x, "i"]
     if weights is not None:
@@ -655,32 +654,27 @@ def bincount(x, weights=None, minlength=0, split_every=None):
         meta = array_safe(np.bincount([]), x._meta)

     if minlength == 0:
-        output_size = np.nan
+        output_size = (np.nan,)
     else:
-        output_size = minlength
+        output_size = (minlength,)

     chunked_counts = blockwise(
         partial(np.bincount, minlength=minlength), "i", *args, token=token, meta=meta
     )
-    chunked_counts._chunks = tuple(
-        (output_size,) * len(c) if i in axis else c
-        for i, c in enumerate(chunked_counts.chunks)
-    )
+    chunked_counts._chunks = (output_size * len(chunked_counts.chunks[0]), *chunked_counts.chunks[1:])

     from .reductions import _tree_reduce

     output = _tree_reduce(
         chunked_counts,
         aggregate=partial(_bincount_agg, dtype=meta.dtype),
-        axis=axis,
+        axis=(0,),
         keepdims=True,
         dtype=meta.dtype,
         split_every=split_every,
         concatenate=False,
     )
-    output._chunks = tuple(
-        (output_size,) if i in axis else c for i, c in enumerate(chunked_counts.chunks)
-    )
+    output._chunks = (output_size, *chunked_counts.chunks[1:])
     output._meta = meta
     return output

[GenevieveBuckley]

Thank you for the feedback @pentschev and @jrbourbeau
Those are very helpful suggestions, and I've updated this PR to include them.

[pentschev]

LGTM, thanks @GenevieveBuckley ! 🙂

[jrbourbeau]

Thanks for all your work on this @GenevieveBuckley and for reviewing @pentschev! This is in

[jrbourbeau]

The fact that we now need NumPy>=1.20

@pentschev could you clarify when NumPy>=1.20 is needed? I noticed we're only skipping test_bincount in the CuPy tests

[GenevieveBuckley]

The fact that we now need NumPy>=1.20

@pentschev could you clarify when NumPy>=1.20 is needed? I noticed we're only skipping test_bincount in the CuPy tests

@jrbourbeau I hope this helps explain it:

If we use array_safe in a function and the cupy test uses assert_eq, then numpy<=1.20 is required for cupy users. Regular users of Dask are ok with any numpy version.

• If numpy>=1.20 then the like= keyword argument is used in the array creation, so everything works properly for everyone (cupy and numpy Dask users alike). However this keyword argument is only available from numpy 1.20 onwards.
• If numpy<1.20 then array_safe ensures that dask doesn't create an error by falling back to using array creation without the like= keyword argument. That means we get a numpy array by default instead of a cupy one. That's no good for cupy users, but it works just fine for people using Dask with numpy array chunks. That's why we need to skip the cupy test is the numpy version is too low.

[jrbourbeau]

Great, thanks for that very clear explanation @GenevieveBuckley