PyTorch - Efficient way to apply different functions to different 'row/column' of a tensor

user_15095194 0
Viewed 366

Let's say I have a 2-d tensor:

x = torch.Tensor([[1, 2], [3, 4]])

Is there an efficient way to apply one function to the first 'row' [1, 2] and apply a second different function to the second row [3, 4]? (Doesn't have to be a row, could be across any dimension)

At the moment, I use the following code: Say I have my two functions, f and g, for example,

def f(z):
    return 2 * z
def g(z):
    return 0.5 * z

Then, to apply them to seperate rows I would do:

torch.cat([f(x[0]).unsqueeze(0), g(x[1]).unsqueeze(0)], dim = 0)

which gives the desired tensor [[2, 4], [1.5, 2]].

Obviously, in this 2-d example this solution is fine, but it seems a bit clunky. Is there a better way of doing this? Particularly in higher dimensions or when there are a large number of elements in the chosen dimension

pytorch python
1 Answers

A handy tip is to slice instead of selecting to avoid the unsqueeze step. Indeed, notice how x[:1] keeps the indexed dimension compared to x[0].

This way you can perform the desired operation in a slightly shorter form:

>>> torch.vstack((f(x[:1]), g(x[1:])))

Optionally you can use vstack to not have to provide dim=0 to torch.stack.

Alternatively, you can use a helper function that will apply both f and g:

 >>> fn = lambda a,b: (f(a), g(b))

And split the tensor inline with torch.Tensor.split:

 >>> torch.vstack(fn(*x.split(1)))
user_63313690
Related