import torch x = torch.ones(3, 3) factors = [lambda x: 2*x, lambda x: 3*x, lambda x: 4*x] indices = torch.tensor([0, 1, 2]) def multiply_row_by_factor(row, idx): return factors[idx](row) result = torch.vmap(multiply_row_by_factor, in_dims=(0, 0))(x, indices) # Original Tensor # tensor([[1., 1., 1.], # [1., 1., 1.], # [1., 1., 1.]]) # Desired Result # tensor([[2., 2., 2.], # [3., 3., 3.], # [4., 4., 4.]]) As the title states, I am looking for a way to call multiple functions on multiple rows of a tensor. I show a minimally reproducible example for simplicity. I am aware that vmap is only meant to be called with one function. I am just using it here as an example to communicate what I am trying to do. This particular approach doesn't work due to the fact that idx is a BatchedTensor. The functions here are lambdas, but in reality, my functions are composed of complex transformations that I would prefer not to decompose to get this to work. Is there any way to achieve something like this? Something cleaner than pytorch streams? Continue reading...
