How can I do these operations in one call?

C++
1

mask has shape [125]
dots has shape [100, 125]

First I want to multiply each element in column c in dots by mask[c]

Second I want to sum each column c of dots to end up with a tensor of shape [100]

        for (int i = 0; i < 125; i += 1)
        {
            for (int j = 0; j < 100; j+= 1)
            {
                dots.index({j, i}) = dots.index({j, i}) * mask.index({i});
            }
        }

        torch::Tensor final_add = torch::zeros({100}, torch::dtype(torch::kFloat64).requires_grad(false));

        for (int i = 0; i < 125; i += 1)
        {
            for (int j = 0; j < 100; j += 1)
            {
                final_add.index({j}) = final_add.index({j}) + dots.index({j, i});
            }
        }
2

You could use broadcasting for the first operation and sum(1) for the second one as seen here:

# setup
mask_ref = torch.randn(125)
dots_ref = torch.randn(100, 125)

# loop approach
mask = mask_ref.clone()
dots = dots_ref.clone()

for i in range(125):
    for j in range(100):
        dots[j, i] = dots[j, i] * mask[i]
    

final_add = torch.zeros(100)

for i in range(125):
    for j in range(100):
        final_add[j] = final_add[j] + dots[j, i];
    
# broadcasting
mask = mask_ref.clone()
dots = dots_ref.clone()

dots = dots * mask.unsqueeze(0)
result = dots.sum(1)

print((result - final_add).abs().max())
> tensor(7.6294e-06)
1 Like