I have two tensors, A and B. A is of dimension [n+1, m]. B is of dimension [n,k] and all value are from -1 to m-1. What I need to to is check if the value in if any value in B is equal to -1. If it is, I will set it to value A[n+1]. If it is not, I will set the value to A[i]. For example, let A = [x_0, x_1, x_2, x_3,x_4] and B = [[0,-1,-1],[1,2,-1],[2,0,1]]. Then, I want to end up with C = [[x_0,x_4,x_4],[x_1,x_1,x_4],[x_3,x_3,x_3]]. I want this to be as efficient as possible, so I’m trying to avoid loops and stick to torch functions.

Any suggestion?

I don’t quite understand the desired output, as it seems you are not setting the values in `C`

to `A[i]`

if `B[i]!=-1`

.

E.g. the second row is given as `[x_1, x_1, x_4]`

and the last one as `[x_3, x_3, x_3]`

. Could you explain a bit more how these results are created?

I think this summarizes it well. If B[i][j]!=-1, then C[i][j]=A[i]. else C[i][j]=A[-1]. I tried using where, but An issue I’m running into is A[i][j].size() != B[i][j].size(). I should specify, x_i is a tensor of length > 1

If it helps, this loop does what I want, but I feel like there is a pytorch way to do what I want without a loop:

C = []

for x in range(len(B)):

CPrime= []

for y in range(len(B[x])):

if B[x][y]!=-1:

CPrime.append(A[x])

else:

CPrime.append(A[-1])

C.append(torch.stack(CPrime))

C = torch.stack(C)

(Sorry for formatting, I’m new to the forums)

where(B!=1, A, A[-1]) would be what I want, but this doesnt work because A is 3 dimensional and B is two dimensional

This should work:

```
A = torch.arange(5)
B = torch.tensor([[0,-1,-1],[1,2,-1],[2,0,1]])
C = []
for x in range(len(B)):
CPrime= []
for y in range(len(B[x])):
if B[x][y]!=-1:
CPrime.append(A[x])
else:
CPrime.append(A[-1])
C.append(torch.stack(CPrime))
C = torch.stack(C)
print(C)
D = A[torch.arange(B.size(0))].unsqueeze(1).expand_as(B).clone()
D[B==-1] = A[-1]
print((C==D).all())
> tensor(True)
```

You can post code snippets by wrapping them into three backticks ```

1 Like

Thanks so much. I had to make a slight change since A is multidimensional. Here is what I ended up with:

```
D = A[torch.arange(B.size(0))].unsqueeze(1).expand(-1, B.size(1),-1).clone()
D[B==-1]=torch.full_like(A[0], -float("inf"))
```