# Expand a tensor, varying the expansion dimension by row

I would like to expand a 2d tensor to a 3d tensor, varying which dimension is expanded in each row of the tensor. For example, I’d like to take this 2d tensor:

``````values = torch.tensor([[1, 2],
[3, 4],
[5, 6],
[7, 8],
[9, 10]])
``````

And transform it to this 3d tensor:

``````tensor([[[ 1,  1],
[ 2,  2]],

[[ 3,  4],
[ 3,  4]],

[[ 5,  5],
[ 6,  6]],

[[ 7,  7],
[ 8,  8]],

[[ 9, 10],
[ 9, 10]]])
``````

where the dimension to expand for each row is given by these indices:

``````torch.tensor([0, 1, 0, 0, 1])
``````

The following code produces the desired result, but is inefficient:

``````def expand_slow(expand_indices, values):
expanded_values = torch.stack([
row_val.expand([2,-1]).transpose(expand_indices[row_idx], 1)
for row_idx, row_val in enumerate(torch.unbind(values, dim=0))
], dim=0)
return expanded_values

expand_indices = torch.tensor([0, 1, 0, 0, 1])

values = torch.tensor([[1, 2],
[3, 4],
[5, 6],
[7, 8],
[9, 10]])

expand_slow(expand_indices, values)
``````

What’s the best way to perform this operation using efficient pytorch functions? Thank you!