I'm trying to get and set indices in a particular tensor dimension without reshaping if possible. I've been able to find the torch.index_select function which does what I want when getting values, but I haven't found an analogous setter function yet. Does one exist?
For context I have a tensor and a set of indexes
class_energy = torch.rand(3, 10, 32, 32)
class_logits = torch.empty_like(class_energy)
idxs = [2, 3, 5, 7]
I want to access items at those indices in a particular dimension so I can perform a log_softmax.
If I knew dim a-priori then I could simply use the fancy __getitem__ / __setitem__ syntax: For example, if dim=1, then class_energy[:, idxs]. Likewise, if dim=2 -> class_energy[:, :, idxs], dim=0 -> class_energy[idxs], etc...
In the case where dim=1, I essentially want this:
class_logits[:, idxs] = F.log_softmax(class_energy[:, idxs], dim=1)
Unfortunately, I don't know the value of dim ahead of time. Sure I could build up the fancy index ahead of time via:
fancy_index = tuple([slice(None)] * dim + [idxs])
class_logits[fancy_index] = F.log_softmax(class_energy[fancy_index], dim=dim)
However, I'm wondering if there is a better way to do this. For the case of __getitem__, I know for a fact that there is. The following code using torch.index_select is equivalent
fancy_index = tuple([slice(None)] * dim + [idxs])
index = torch.LongTensor(idxs).to(class_energy.device)
class_logits[fancy_index] = F.log_softmax(torch.index_select(class_energy, dim=dim, index=index, dim=dim))
Not only is it functionally the same, index_select is much faster (I've seen 2x improvement) than using the fancy getitem syntax.
My issue is that I can't seem to similar functionality for the __setitem__ portion of the code. It would be really nice if I could get rid of the fancy index all together. I've looked into Tensor.put_, torch.index_put, and torch.select, but none of these seem to have the functionality I want. Is there something I'm missing, or is the fancy index the only way to solve this problem currently?