Assignment to Tensor in Similar Way to Torch.select

1

Is there an easy way to assign values to a tensor in a way that’s similar to torch.select?

For instance, say I have a tensor with n dimensions and I want to change the mth “row” of the kth dimension to be 1. If n = 5, m = 2, and k = 3, I could do something like tensor[:, :, :, 2, :] = 1 - but let’s say that m, n, and k are not known beforehand, or I want to loop through different values of m and k, etc. If I just wanted to retrieve the values of these slices, torch.select(tensor, 3, 2) would work, but I want to be able to make changes to this slice.

2

Could you post a (slow) example how exactly this assignment should look like which we could use as a reference?

3

Sure - the exact use case I’m thinking of looks like this, except the size of the tensor can be different in different cases.

import torch
tensor = torch.ones([5,6,7])

tensor[0] *= 0.5
tensor[-1] *= 0.5
tensor[:, 0] *= 0.5
tensor[:, -1] *= 0.5
tensor[:, 0] *= 0.5
tensor[:, -1] *= 0.5
tensor[..., 0] *= 0.5
tensor[..., -1] *= 0.5

I’m hoping there’s a way to somehow loop through the dimensions, sort of like

for dim in range(len(tensor.shape)):
    torch.select(tensor, dim, 0) *= 0.5 # But with a different function than torch.select
    torch.select(tensor, dim, -1) *= 0.5