There’s a way to do this if we restrict your mask M:
Let’s say you have a mask M such that no two 1’s in the mask are in the same row or column.
Then the following should equivalent to (z @ y) * M, where the @ sign is matrix multiplication:
(z.t() * (y @ M.t())).to_dense().sum(dim=0).diag() @ M
I think pytorch does support sparse x dense -> sparse via torch.mm. If you need a dense x sparse -> sparse (because M will probably be sparse), you can use the identity AB = ( AB )^T ^T = (B^T A^T)^T.
I haven’t actually tried out the expression above for your use case, though.