Inserting part of tensor into another

1

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?

2

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?

3

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

4

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

5

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 ``` :wink:

1 Like
6

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"))