Fastest way to implement

Hi,

I want to fill a tensor like this:

gainW = torch.tensor(1.1)
gainH = torch.tensor(2.5)
res = 5
a = torch.zeros(res,res)
a[0,0] =0.1
for i in range(res):
    for j in range(res):
        if i==0 and j==0:
            continue
        if j==0:
            a[i,j] = torch.cos(a[i-1,0]*gainH)
        else:
            a[i,j] = torch.sin(a[i,j-1]+gainW)

print(a)
tensor([[ 0.1000,  0.9320,  0.8955,  0.9112,  0.9046],
        [ 0.9689,  0.8785,  0.9180,  0.9016,  0.9086],
        [-0.7523,  0.3408,  0.9916,  0.8674,  0.9224],
        [-0.3049,  0.7139,  0.9706,  0.8777,  0.9184],
        [ 0.7233,  0.9683,  0.8788,  0.9179,  0.9017]])

But this is extremely inefficient, How can I improve it? Any ideas would be appreciated!

BTW, I want to calculate later the gradient with respect to gainW,gainH!

import torch

gainW = torch.tensor(1.1)
gainH = torch.tensor(2.5)
res = 5
a = torch.zeros(res,res)
a[0,0] = 0.1
for i in range(1, res):
    a[i, 0] = torch.cos(a[i-1, 0] * gainH)

for j in range(1, res):
    a[:, j] = torch.sin(a[:, j-1] + gainW)

This might be faster.
As your case requires the previous computation result to get the next, it’s difficult to avoid for loops entirely.

1 Like

Thnx, Do you think this is the best I can do? Do you think putting it in JIT (trace) would help? Does trace keep the gradient flow?