Hi Wasabi!



 wasabi:

is there a way to naturally broadcast P?


No.  Quoting from the documentation for torch.gather() 5:

Note that input and index do not broadcast against each other.

(Also, I’m not aware of any way to do what you want with pytorch
indexing operations other than gather(), although I may have
overlooked something.)
Best.
K. Frank