seba-1511:

By spelling of the contraction, do you mean writing the operation as an einsum ?


No, einsum will itself use bmm, I thought of materializing the elementwise product and .sum.
(I do have a branch somewhere that uses TensorIterators for einsum instead. It’s so terrible on CPU (no AVX) that I didn’t look at GPU, but if you want to benchmark it on GPU, I can push it.)
Best regards
Thomas