is there any way to use softmas as aggregation for sparse.mm?

[businessabstraction]

Hi,

I am trying to use sparse.mm for a Graph NN, representing adjacency matrix as a sparse tensor. However, max doesn't work well with Autograd (and is not implemented for sparse tensors anyway). Any suggestion on how I can combine sparse trensor with softmax?

Thank you