Is there a function to divide a n dimensional tensor with an n-1 dimensional tensor

Viewed 66

I am not sure how to word this question properly so I will show some examples to describe the desired behaviour.

I'm looking to divide a tensor in this specific way.

  1. Divide a vector by 1 scalar, for example [1, 2, 3, 4, 5] divided by 2 = [0.5, 1, 1.5, 2, 2.5]

  2. Divide a matrix by 2 scalars, for example [[1, 2, 3], [2, 3, 4]] by [2, 4] = [[0.5, 1, 1.5], [0.5, 0.75, 1]]

  3. Divide a 3 dimensional tensor by a 2 dimensional tensor, for example [[[1, 2, 3], [2, 3, 4]], [[4,5,6], [7,8,9]], [[9,8,6], [1,2,3]]] divided by [[1, 2], [3, 4], [5, 6]] = [[[1, 2, 3], [1, 1.5, 2]], [[4/3,5/3,2], [7/4,2,9/4]], [[9/5,8/5,6/5], [1/6,2/6,3/6]]]

  4. Divide a N dimensional tensor by a N-1 dimensional tensor...

I'm looking for a pytorch way to do this.

1 Answers

You can expand the dimensions of the N-1 dimensional tensor to make it broadcastable with the N-dimensional tensor.

tensor_a / tensor_b.unsqueeze(-1)

This generalizes, even when the denominator is a scalar. The -1 dimension means the last dimension. This follows from python indexing rules, in which sequence[-1] gives you the last element of the sequence.

a = torch.as_tensor([1, 2, 3, 4, 5])
b = torch.as_tensor(2)
a / b.unsqueeze(-1)
# tensor([0.5000, 1.0000, 1.5000, 2.0000, 2.5000])

a = torch.as_tensor([[1, 2, 3], [2, 3, 4]])
b = torch.as_tensor([2, 4])
a / b.unsqueeze(-1) 
# tensor([[0.5000, 1.0000, 1.5000],
#         [0.5000, 0.7500, 1.0000]])

a = torch.as_tensor([[[1, 2, 3], [2, 3, 4]], [[4,5,6], [7,8,9]], [[9,8,6], [1,2,3]]])
b = torch.as_tensor([[1, 2], [3, 4], [5, 6]])
a / b.unsqueeze(-1) 
# tensor([[[1.0000, 2.0000, 3.0000],
#          [1.0000, 1.5000, 2.0000]],
#
#         [[1.3333, 1.6667, 2.0000],
#          [1.7500, 2.0000, 2.2500]],
#
#         [[1.8000, 1.6000, 1.2000],
#         [0.1667, 0.3333, 0.5000]]])
Related