Add NumPy's kron for Dask Arrays

[jakirkham]

Would be nice to have a Dask Array implementation of NumPy's kron. A generalization of outer products for two tensors.

[jakirkham]

NumPy implements kron using outer and reshaping. Would be possible to do the same thing for Dask Arrays. That said, expect there are better options. Suggestions welcome.

[mrocklin]

Is this doable with atop or slicing syntax do you think?

Marking this as good first issue, but it may not be. I think that this requires someone to think about it a bit, and then look at the dask/array/core.py::atop function and see if there is a nice way to implement it.

[jakirkham]

Avoiding reshaping is probably better. It may be possible with atop, but agree it requires a bit of thought.

[mathdugre]

Hi, is this issue solved? I would like to work on it

[jakirkham]

It is still open. Please feel free to work on it if it is of interest. Would be happy to review a PR. 😄

[Madhu94]

@mathdugre Please do let me know if you are working on this, I'd like to give it a try

[mathdugre]

@Madhu94 I didn't get the time to work on it, feel free to do so :)

[Madhu94]

thanks! Starting on it.

[gitkarma]

@Madhu94 any progress on it. I was also looking for the feature

[Madhu94]

@gitkarma No I could not find the time. Still interested in exploring this, but if anyone else wants to push this through in the interests of time, please take it up

[eric-czech]

Is this rechunking into 1x1 chunks for one of the arguments likely to be as bad or worse than reshaping? I can't figure out how to do it otherwise with blockwise, though this seems like it could be on the right track:

x = da.arange(36).reshape(9, 4).rechunk((1, 1))
y = da.ones(36).reshape(9, 4).rechunk((3, 2))
z = da.blockwise(
    np.multiply, 'ij', x, 'ij', y, 'xy', concatenate=True, dtype='f8',
    adjust_chunks={'i': y.shape[0], 'j': y.shape[1]}
)
z.shape # -> (81, 16)
# Compute on kron inputs to silence "Not implemented by Dask ..." warning
assert np.all(np.kron(x.compute(), y.compute()) == z.compute())
z.compute()
array([
       [ 0.,  0.,  0., ...,  3.,  3.,  3.],
       [ 0.,  0.,  0., ...,  3.,  3.,  3.],
       [ 0.,  0.,  0., ...,  3.,  3.,  3.],
       ...,
       [32., 32., 32., ..., 35., 35., 35.],
       [32., 32., 32., ..., 35., 35., 35.],
       [32., 32., 32., ..., 35., 35., 35.]
])

Separately, is this a bug or expected behavior?

x = da.arange(36).reshape(9, 4).rechunk((1, 1))
y = da.ones(36).reshape(9, 4).rechunk((3, 2))
z = da.blockwise(
    np.multiply, 'ij', x, 'ij', y, 'xy', concatenate=True, dtype='f8',
    # Leave this out to show that inferred shape is wrong
    # adjust_chunks={'i': x.shape[0], 'j': y.shape[1]}
)
print(z.shape) # --> (9, 4) 
print(z.compute().shape) # --> (81, 16)
# And this still works:
assert np.all(np.kron(x.compute(), y.compute()) == z.compute())

[jakirkham]

Thanks for sharing Eric! 😄

Yeah that's a good question. Rechunking is generally expensive. However reshaping will result in some smearing and remixing of chunks, which comes with its own cost. Probably requires some benchmarking to know for sure.

That said, there might be ways to tweak what you came up with to improve things further before measuring timings. 🙂

[eric-czech]

This might be an improvement:

x = da.arange(36).reshape(9, 4).rechunk((3, 2))
y = da.ones(36).reshape(9, 4).rechunk((3, 2))
z = da.blockwise(
    np.kron, 'ab', x, 'ij', y, 'ij', concatenate=True, dtype='f8',
    new_axes={'a': x.shape[0] * y.shape[0], 'b': x.shape[1] * y.shape[1]}
)
assert np.all(np.kron(x, y) == z.compute())
assert z.shape == z.compute().shape
z.compute()
array([[ 0.,  0.,  0., ...,  3.,  3.,  3.],
       [ 0.,  0.,  0., ...,  3.,  3.,  3.],
       [ 0.,  0.,  0., ...,  3.,  3.,  3.],
       ...,
       [32., 32., 32., ..., 35., 35., 35.],
       [32., 32., 32., ..., 35., 35., 35.],
       [32., 32., 32., ..., 35., 35., 35.]])

I don't think it's what I want though since passing concatenate=False gives blocks in lists and building one giant result as the concatenation of all kronecker products would be a bad idea. I'll keep experimenting ..