Support for np.take_along_axis?

[phizaz]

Numpy take_along_axis reference: https://docs.scipy.org/doc/numpy/reference/generated/numpy.take_along_axis.html

I have found a use for zarr in reinforcement learning which needs to keep a good size of array in memory. Fortunately, this array is very compressible. I think it could make a good use of a library like zarr. For some reason, take_along_axis is important for generalization of code to larger settings. I am here to request such feature.

Thank you.

[jakirkham]

Have you taken a look at Dask? Many users of Zarr use it in conjunction with Dask to have efficient array storage in conjunction with a NumPy-like API.

[phizaz]

@jakirkham I have taken a look, but it seems unclear to me how can I use zarr as underlying storage for dask. Would be much appreciated if you point out a direction!

[jakirkham]

If you have an existing Zarr Array (z), you can create a Dask Array (d) like so.

import dask.array as da
d = da.from_array(z, chunks=z.chunks)

[phizaz]

@jakirkham It seems that take_along_axis is not supported by dask at the moment. Refer: dask/dask#3663.