Your PR is an impressive piece of code @ravwojdyla !

I know this PR is closed. But I'd appreciate it if you would answer one last question from me.

This is the only line that puzzles me — while I understand why the number of the dimensions of the output chunk should match what blockwise expects, I can't figure out why the new axis is inserted in the last position instead of the second to last position. Perhaps it doesn't matter in the end, but I'd like to know why.

I mean, why not the following:

    # to add one dummy dimension back
    return chunk[..., np.newaxis, :]

which more closely matches the expected output of the blockwise call?

Thanks!

👋 @ParticularMiner my initial reaction would be that it probably doesn't matter (tho might require some changes it in the downstream logic from blockwise). That extra dummy dimension is "squeezed" out later afair, and as you have pointed out - it's there to make blockwise happy at the metadata level. That said, it's been a while, and it's just my initial reaction.

Thank you very much for your reply @ravwojdyla despite this conversation being almost a year old!

Indeed, it turns out through testing that it doesn’t matter, which puzzled me.

Also the downstream logic after blockwise operates on the assumption that the contraction-axis is the second-to-last axis of the output of blockwise.

@ravwojdyla

The following is not vital (performance-wise). But just to make you aware:

While mirroring your code elsewhere, it turned out that by using

    # to add one dummy dimension back
    return chunk[..., np.newaxis, :]

instead of

    # to add one dummy dimension back
    return chunk[..., np.newaxis]

one can avoid the extra conditional branches in the downstream logic after blockwise by replacing them with a single sum. That is, this line:

    out = out.sum(axis=-2)

can replace the following lines:

    # When we perform reduction, we need to worry about the last 2 dimensions
    # which hold the matrices, some care is required to handle chunking in
    # that space.
    contraction_dimension_is_chunked = (
        max(min(a.chunks[-1], b.chunks[-2])) < a.shape[-1]
    )
    b_last_dim_max_chunk = max(b.chunks[-1])
    if contraction_dimension_is_chunked or b_last_dim_max_chunk < b.shape[-1]:
        if b_last_dim_max_chunk > 1:
            # This is the case when both contraction and last dimension axes
            # are chunked
            out = out.reshape(out.shape[:-1] + (1, -1))
            out = out.sum(axis=-3)
            out = out.reshape(out.shape[:-2] + (b.shape[-1],))
        else:
            # Contraction axis is chunked
            out = out.sum(axis=-2)
    else:
        # Neither contraction nor last dimension axes are chunked, we
        # remove the dummy dimension without reduction
        out = out.reshape(out.shape[:-2] + (b.shape[-1],))

I haven't tested this yet for dask itself. But it works in my particular application. If you feel this is worth exploring for dask, then just drop me a line, and I'll create a PR with extensive testing. But I doubt it will significantly improve current performance. Rather it will merely create more readable code. Which is often useful. 😉

@ParticularMiner that's great! Given that those reshape ops should be largely metadata operation, I agree that I won't expect performance impact, but I do love the simpler code. Definitely +1 to PR, iff you test that it works well for diff chunking schemes (as mentioned in the comment above).

@ravwojdyla

Many thanks for your interest!

I've made the suggested changes on my local git and the relevant existing unit-tests (in test_matmul()) have all passed.

The only thing left now is to run your benchmarking jupyter notebook above. (IMO, the results there will be similar. Still it is always safer to check.) But to do that, I need access to a computing cluster, which I currently do not have. Any pointers on where I can access one (preferably for free)? 😉

@ravwojdyla

Sorry. I misunderstood. It seems you ran your notebook on a standard laptop, not a multi-node cluster. I can probably do that too.

Hi @ravwojdyla

You’ll find my suggested PR at #8423.

Currently the CI tests are broken for other reasons. So it is not clear that the PR checks out.