Flattening Dask Arrays without spreading out chunks

[jakirkham]

Currently functions like ravel and flatten follow the same traversal that one would expect when using NumPy. This makes sense for users coming from NumPy, who would expect this. Also there are some Dask Array functions that rely on this behavior to match their NumPy counterparts. However this means that ravel and flatten must spread out the chunks (rechunk the array), which can have a non-negligible performance cost associated.

That said, not all cases rely on matching the same order that NumPy would provide when flattening out an array. For cases like this, a nice alternative would be to flatten chunks themselves and merely stitch together the flattened chunks into a new 1-D array. This strategy would require no rechunking and would be embarrassingly parallel. Thus it would avoid the performance penalties that ravel and flatten have today.

Assuming this strategy is reasonable, there are a few ways we could go about implementing it.

1. Allow some additional options to the order parameter to handle this need.
2. Add a new parameter to toggle NumPy or chunk-based traversal strategies.
3. Include a config option to enable this behavior for Dask Arrays more generally.
4. Add a new function entirely for this behavior.
5. ?

Thoughts on this?

[jakirkham]

cc @shoyer @pentschev @GenevieveBuckley (to get a few different perspectives on this)

[shoyer]

Yes, this would be useful. I think either a new option in the order parameter or a new parameter to these functions/method would work.

[jakirkham]

Just to provide more context for item 3 (config option), this could be useful when the user does not have direct access to ravel/flatten, but still wants this behavior to be used under the hood. For example, argwhere flattens the array and nonzero also does so indirectly.

[shoyer]

I would much rather thread a new argument through argwhere/nonzero in dask than add a configuration option. Global configuration options lead to spooky "action at a distance" that can be very hard to debug, e.g., if unexpected operations also happen to indirectly use flattening.

[jakirkham]

Good point. Just wanted to make sure it made sense why that was in the list. 🙂

[jakirkham]

FWIW here is the code I hacked together to do this myself. Not to say it should be implemented exactly like this, but hopefully it is useful when this actually gets implemented somewhere.

from itertools import product, starmap
from operator import methodcaller

import numpy as np

from dask.array import Array
from dask.base import tokenize
from dask.core import flatten
from dask.highlevelgraph import HighLevelGraph


def ravel_chunks(arr):
    name = "ravel_chunks-" + tokenize(arr)
    chunks = (tuple(map(np.product, product(*a.chunks))),)
    dsk = {(name, i): (methodcaller("ravel"), k)
           for i, k in enumerate(flatten(arr.__dask_keys__()))}
    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[arr])
    res = Array(graph, name, chunks, arr.dtype)
    
    return res

[pentschev]

Sorry for joining the discussion late, but while I agree with @shoyer's point on global configurations, I'm also growing more concerned about non-NumPy API arguments. If someone wants to use Dask today calling np.ravel(dask_array, non_numpy_arg=some_value), non_numpy_arg would have to be removed for a non-Dask array in the future.

[jakirkham]

Does this mean you would prefer option 1 or 4?

[pentschev]

Actually, I can't think of any good solutions that won't introduce issues on one way or another. Option 1 would probably introduce the same problem as I've mentioned before. Option 4 doesn't do that, but on the other hand, it would introduce another function, that wouldn't work with __array_function__. I think option 4 would at least be more explicit to users, since they won't need to check which library introduces some non-NumPy argument, but not necessarily a better solution than the others.

How often do such cases happen? If not very often, I would prefer, by order of highest preference, solutions 4 or 3.

[shoyer]

I have no strong preference between options (1), (2) and (4).

[GenevieveBuckley]

The major point I want to make is that if this is implemented, an explanation of it should also go in the 'best practices for dask arrays' docs - I bet there are users out there who'd be grateful to know about the performance implications.

[GenevieveBuckley]

For the question you actually asked:

1. I'd also vote no for option 3 (I agree with @shoyer on this, he explains the potential pitfalls well upthread).
2. I'd want to prioritise an API that is compatible with numpy as far as possible. Ideally it would be relatively easy for users to take code they wrote for dask arrays and re-use it for numpy arrays without it breaking. In cases where that's not possible, maybe it's best to try and make things obviously non-compatible with numpy? Would that help narrow down the remaining options?

My thinking around point 2 is very fuzzy, so I'd welcome anyone else expanding on these points or arguing against them.