CODEWITHSUNDEEP
pythonpytorchvectorization
DimSum
DimSum

Reputation: 123

Pairwise difference between tensors with PyTorch

I have the following problem: given two tensors X of shape (T,n,d) and Y of shape (R,m,d), I want to compute the tensor D of shape (T,R,n,m,d) such that, for all 0 <= t < T, 1 <= r < R, 1 <= i < n, 1 <= j < m,

D[t,r,i,j] = X[t,i] - Y[r,j]

Can we compute D with Pytorch without using any loop?

I know that when given two tensors X of shape (n,d) and Y of shape (m,d) we can compute the tensor D of shape (n,m,d) such that, for all 1 <= i < n, 1 <= j < m, D[i,j] = X[i] - Y[j] using

x.unsqueeze(1) - y

What I am looking for is a similar trick for the initial problem.

Upvotes: 2

Views: 632

Answers (1)

trialNerror
trialNerror

Reputation: 3563

I think the following snippet might achieve what you need :

# reshapes X to (T, R, n, m, d)
X_rs = X.view(T, 1, n, 1, d).expand(-1, R, -1, m, -1) 
# reshapes Y to (T, R, n, m, d)
Y_rs = Y.view(1, R, 1, m, d).expand(T, -1, n, -1, -1)
# Compute the coefficient-wise difference
D = X_rs - Y_rs

Note that the expand operation does not allocate new memory. It's just the generalization of the "trick" you mentioned

Upvotes: 1

Related Questions