Search code examples
pythonpytorchtorch

PyTorch Column-wise reduction using multiplication

I'd like to reduce columns from a Torch tensor by multiplying all values from the same row. So, for instance:

x = torch.tensor([[1,1,1],[1,1,0],[1,1,2], [1,2,2]])

Shape is 4*3. After the reduction, I'd like to have a tensor of shape 4 with each value being the product of each column ie.

x_reduced = torch.tensor([1,0,2,4])

Is there a torch operator to do that easily?

Solution

  • Yep, the function call is simple: torch.prod(x,dim = 1).