Can one of the admins verify this patch?

Admins can comment ok to test to allow this one PR to run or add to allowlist to allow all future PRs from the same author to run.

edit: added benchmark for sparse implementation.
@zklaus, there are a few deviation from your original code in climix:

• I removed sparse arrays as it's not a dependency for dask.
  Performance wise, on my laptop, it looks better without sparse.

from dask_take_along_axis import dask_take_along_axis # copy-pasted module from climix

%timeit dask_take_along_axis(dask_random_arr, from_array(top50_np), axis=0).compute()
6.33 s ± 60.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

• I remove the check of shapes because it was breaking a use case. In particular, with this assert it break the following test dask/array/tests/test_slicing.py::test_take_along_axis__indexing_twice_same_1darray
• I removed the meta argument passed in blockwise as it was a sparse array. I suspect it some sort of optimization to give blockwise what is the expected shape but it's not really documented.
• Also some variable renaming and added types, docs and integration tests.

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

     15 files  ± 0       15 suites  ±0   3h 29m 26s ⏱️ + 7m 10s
 13 124 tests + 3   12 185 ✅ + 3     931 💤 ±0  8 ❌ ±0 
162 507 runs  +39  142 400 ✅ +39  20 099 💤 ±0  8 ❌ ±0 

For more details on these failures, see this check.

Results for commit 105ef10. ± Comparison against base commit dafb6ac.

Is there someone available to review/comment on this PR ?

Sorry that I don't have the bandwidth to interact properly :/ One thing that concerns me a little is that the test/benchmark isn't really testing much dask because it is conducted with a dask array that consists of a single numpy array as it's only chunk. In that case, performance would probably be expected to be slightly worse than numpy, but it's also not really representative.

Particularly, the point of the sparse array, iirc, was that there is one array for every chunk that covers all the chunks. If this is sparse, that's ok (or at least I thought so when writing it), but it is prohibitive if dense, so I don't think replacing the use of sparse with dense matrices is the way to go. If we want to avoid the use of sparse as a dependency, the aggregation step needs to be re-thought.

As a practical step forward, @bzah, could you do a larger test? Why not apply this to some larger climate data in at least a local cluster with several workers? I suspect that perhaps more than compute time, memory requirements without the use of sparse would balloon.

@quasiben maybe this is something your team could review?

Related to this PR is a proposal to add take_along_axis to the Array API Standard: data-apis/array-api#808.

If there are concerns regarding adding take_along_axis to Dask, it would be great to hear those concerns so that we can take those into account when formalizing expected behavior.