I have an operations like this:
x = torch.randint(16, size=(128, 1, 3, 32, 32), dtype=torch.float32)
h = torch.randint(16, size=(1, 6, 3, 32, 32), dtype=torch.float32)
z = (x*h).sum(2) # output shape (128, 6, 32, 32)

Is there a way to avoid the memory allocation of the intermediary step x*h? A mix of broadcast, inplace and reduction on a single call?