Express matmul vjps in terms of matmul

[j-towns]

No description provided.

[duvenaud]

Nice, is the motivation mainly to get rid of the call to einsum?

[j-towns]

Nice, is the motivation mainly to get rid of the call to einsum?

Yes. There should be a performance advantage to this version, particularly for small arrays, by avoiding all the parsing that einsum and its vjp have to do, I haven't profiled it yet though.

[mattjj]

Thanks, Jamie!

I thought matmul might also be faster than einsum because it could more readily call BLAS routines, but after poking around it looks like matmul's implementation kinda sucks and just calls into einsum anyway for the stacked case. I think that's an upstream numpy issue, though, and I would guess this implementation is at least as fast as the previous one.

[j-towns]

Oh that's a shame :(. Even dot sucks in some cases because BLAS doesn't (yet) support arbitrarily strided arrays. There was a good description of how numpy deals with different strides posted here:

how do you perform a GEMM operation on two arrays with arbitrary strides? Currently, NumPy attempts to detect a number of special cases: if the strides in both arrays imply a column-major layout, then call BLAS directly; if one of them has strides corresponding to a row-major layout, then set the corresponding transA/transB argument, etc. – and if all else fails, either copy the data into a contiguous buffer, or else fall back on a naive triple-nested-loop GEMM implementation. (There's also a check where if we can determine through examining the arrays' data pointers and strides that they're actually transposes of each other, then we instead dispatch to SYRK.)