Search code examples
pythonloopspytorchmatrix-multiplication

PyTorch: How to multiply elements in a list containing tuples of integer and tensor


I have the following list ABC_lst containing tuples of an integer and a tensor.

Code:

import numpy as np
import torch

A_int = 40
A_tsr = torch.tensor(np.array([[1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5]]))
A_tpl = (A_int, A_tsr)

B_int = 42
B_tsr = torch.tensor(np.array([[4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8]]))
B_tpl = (B_int, B_tsr)

C_int = 38
C_tsr = torch.tensor(np.array([[7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11]]))
C_tpl = (C_int, C_tsr)

ABC_lst = [A_tpl, B_tpl, C_tpl]
ABC_lst

Output:

[(40, tensor([[1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5]])),
 (42, tensor([[4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8]])),
 (38, tensor([[ 7,  8,  9, 10, 11],
          [ 7,  8,  9, 10, 11],
          [ 7,  8,  9, 10, 11],
          [ 7,  8,  9, 10, 11],
          [ 7,  8,  9, 10, 11]]))]

How do I multiply the integer with the corresponding tensor, for eg. multiply 40 with

tensor([[1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5],
          [1, 2, 3, 4, 5]])

multiply 42 with

tensor([[4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8],
          [4, 5, 6, 7, 8]])

and so on.

The returned object should be a tensor, which looks like this:

tensor([[[ 40.,  80., 120., 160., 200.],
         [ 40.,  80., 120., 160., 200.],
         [ 40.,  80., 120., 160., 200.],
         [ 40.,  80., 120., 160., 200.],
         [ 40.,  80., 120., 160., 200.]],

        [[168., 210., 252., 294., 336.],
         [168., 210., 252., 294., 336.],
         [168., 210., 252., 294., 336.],
         [168., 210., 252., 294., 336.],
         [168., 210., 252., 294., 336.]],

        [[266., 304., 342., 380., 418.],
         [266., 304., 342., 380., 418.],
         [266., 304., 342., 380., 418.],
         [266., 304., 342., 380., 418.],
         [266., 304., 342., 380., 418.]]])

In the above eg., I have 3 "sets" of integer and tensor. How do I generalize a code for the multiplication above for any arbitrary "sets" of integer and tensor?

Would really appreciate it if anyone could help.

EDIT: I need to do all the above in GPU, so need to work with tensors.


Solution

  • Starting from two lists: I the list of integers and X the list of tensors:

    I = [torch.tensor(40), torch.tensor(42), torch.tensor(38)]
    X = [
        torch.tensor([[1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5]]),
        torch.tensor([[4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8], [4,5,6,7,8]]),
        torch.tensor([[7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11], [7,8,9,10,11]]), 
    ]
    

    You can zip both and create a list containing all multiplications results. Then stack this list into a single tensor, like so:

    torch.stack([i*x for i, x in zip(I, X)])
    

    You can, of course, add more elements to your lists.