I would like to create a mechanism for multiplyiing multiple PyTorch Tensors in a single instruction using an Adjacency Matrix. The points are connected to each other in a series, however there is a small catch. Say this was an adjacency matrix denoting the Series connection of the T points:

adj_mat = torch.tensor([
[0, 1, 0, 0, 0],  # T1 connected to T2
[1, 0, 1, 0, 0],  # T2 to T1 and T3
[0, 1, 0, 1, 0],  # T3 to T2 and T4
[0, 0, 1, 0, 1],  # T4  to T3 and T5
[0, 0, 0, 1, 0],  # T5 to T4
], dtype = torch.float32)

I have multiple quantities which will need to be required for the arithmetic operations on the input Tensor:

K = torch.tensor([
    [1.0, 2.0, 3.0, 4.0, 0.0],
    [1.0, 2.0, 3.0, 4.0, 0.0]
], dtype=torch.float32)
C = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype = torch.float32)
S = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype = torch.float32)

val = torch.tensor([[1.0, 2.0, 4.0, 8.0, 16.0],
                      [5.0, 10.0, 15.0, 20.0, 5.0]], dtype = torch.float32)

dev = torch.tensor([11, 30], dtype = torch.float32)

The computations are to be as follows, across each row of the Tensor (wherein each row would be computed in parallel, hopefully):

T1 = T1 + (C(T1) * S(T1) * (dev + (T2-T1) * K(T1))
T2 = T2 + (C(T2) * S(T2) * ((T1-T2) * K(T1) + (T3-T2) * K(T2))
T3 = T3 + (C(T3) * S(T3) * ((T2-T3) * K(T2) + (T4-T3) * K(T3))
T4 = T4 + (C(T4) * S(T4) * ((T3-T4) * K(T3)+ (T5-T4) * K(T4))

wherein the Values of T1-T5 would be sourced from the val array. I tried an approach like so to first check if the weighted differences, i.e., ((T2-T3) * K(T2) + (T4-T3) * K(T3)) are being computed right or not

T_diffs = val.unsqueeze(1) - val.unsqueeze(2)
T_diffs_weighted = T_diffs * adj_mat.unsqueeze(0) * K.unsqueeze(1)
T_diffs_summed = torch.sum(T_diffs_weighted , dim = 2)
print("weighted Differences:", T_diffs_weighted )
print("Summed Differences:", T_diffs_summed )

The output I got is:

weighted Differences: tensor([[[  0.,   2.,   0.,   0.,   0.],
         [ -1.,   0.,   6.,   0.,   0.],
         [ -0.,  -4.,   0.,  16.,   0.],
         [ -0.,  -0., -12.,   0.,   0.],
         [ -0.,  -0.,  -0., -32.,   0.]],

        [[  0.,  10.,   0.,   0.,   0.],
         [ -5.,   0.,  15.,   0.,  -0.],
         [ -0., -10.,   0.,  20.,  -0.],
         [ -0.,  -0., -15.,   0.,  -0.],
         [  0.,   0.,   0.,  60.,   0.]]])

whereas the desired output is¹:

weighted Differences: tensor([[[  0.,   1.,   0.,   0.,   0.],
                                  [ -1.,   0.,   6.,   0.,   0.],
                                  [  0.,  -4.,   0.,  16.,   0.],
                                  [  0.,   0., -12.,   0.,   0.],
                                  [  0.,   0.,   0., -32.,   0.]],

                                 [[  0.,   5.,   0.,   0.,   0.],
                                  [ -5.,   0.,  15.,   0.,  -0.],
                                  [  0., -10.,   0.,  20.,  -0.],
                                  [  0.,   0., -15.,   0.,  -0.],
                                  [  0.,   0.,   0.,  60.,   0.]]])

Summed Differences: tensor([[  1.,   1.,  10., -16.,   0.],
                                 [  0.,   5.,   0.,  65.,   0.]])

The extracted value would be a column, meaning T1 would correspond to the 0th column for example. I cannot seem to multiply the Tensor K correctly. I am able to either of the K-value for an equation right, not both of them.

Finally, any other approaches to execute this multiplication or rather the equation would also be appreciated. The equations are part of a machine learning model which runs on an Nvidia GPU, by the way. I am relatively new with the [un]squeeze() methods of PyTorch which is why this is slightly hard to grasp.

¹The signs of the values in the weighted difference matrices may be inverted as compared to the expected mathematical output from the equation. But the multiplication with the Tensor K is what I am focusing on.