pytorch how to join two tensor like in sql

Sida Zhou 0
Viewed 30

Scenario:

# I start with unsorted, non-consecutive duplicated keys, the order has to be perserved
ordered_keys = torch.tensor([1,5,5,3,1,1]).T # shape 6x1

# blackbox() is a function that only takes non-duplicated array, can be sorted or unsorted
# input shape 3x1
blackbox_input = ordered_keys.unique() # input = [1, 3, 5]

# blackbox() outputs the corresponding features of the keys
# blackbox_output shape 3x2: 3keys, each key has 2 features
blackbox_output = blackbox(blackbox_input) # blackbox_output = [[100,101],[300,301],[500,501]]

# What I want is essentially a table join using torch or numpy,
# such that output shape is 6x2, with the same order as in the beginning, with values:
# [[100,101],
# [500,501],
# [500,501],
# [300,301],
# [100,101],
# [100,101]]

output = ordered_keys.join(blackbox_output) # <--- doesnt work

Question:

How would you do this efficiently within torch or numpy? (avoiding for-loops or pandas or sql etc)

pytorch numpy
1 Answers
ordered_keys = torch.tensor([1,5,5,3,1,1])
# print(ordered_keys)

blackbox_input = torch.tensor([1, 3, 5])
# print(blackbox_input)
blackbox_output = torch.tensor([[100,101],[200,201],[300,301]])
# print(blackbox_output)

print(blackbox_output[ordered_keys.unique(return_inverse=True)[1]])

outputs:

tensor([[100, 101],
        [300, 301],
        [300, 301],
        [200, 201],
        [100, 101],
        [100, 101]])
Sida Zhou0
Related