How to calculate log_softmax for list of tensors without breaking autograd

[InnovArul (Arul)]

@Aeryan

even though I am using the inplace operation += instead of something like torch.sum, torch.sum is called under the hood

Yes, you are right.

@Aeryan

I want to calculate logsumexp for the tensors x = torch.randn(5,3) and y = torch.randn(3, 7)

Can you elaborate a bit more on the context? Why are these tensors of different size?
In your code, it looks like you are considering all the elements of the (possibly) multi-dimensional tensors, as if they are vectors.
So, the below code might work?

x = torch.randn(5,3)
y = torch.randn(5,6)
z = torch.randn(5,9)
composed = torch.cat([x.flatten(), y.flatten(), z.flatten()], dim=-1)
logsumexp = torch.logsumexp(composed, dim=-1, keepdim=True)