Assigning values to Tensor without loops

autograd
1

There must be a better way to implement this function, but I couldn’t find any similar solution. Can somebody help with this? x is of shape (500,6,4).

def myFcn(x):

cos_x2 = torch.cos(x[:,2])
sin_x2 = torch.sin(x[:,2])
cos_x5 = torch.cos(x[:,5])
sin_x5 = torch.sin(x[:,5])
out = torch.ones([x.shape[0], 6, 4])

for i in range(x.shape[0]):
  out[i,:,:] = torch.Tensor([
    [cos_x2[i], 0.0, 0.0, 0.0],
    [sin_x2[i], 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
    [0.0, 0.0, cos_x5[i], 0.0],
    [0.0, 0.0, sin_x5[i], 0.0],
    [0.0, 0.0, 0.0, 1.0]])

return out
2

This seems to be a faster implementation:

def myFcn(x):
    cos_x2 = torch.cos(x[:,2])
    sin_x2 = torch.sin(x[:,2])
    cos_x5 = torch.cos(x[:,5])
    sin_x5 = torch.sin(x[:,5])
    out = torch.ones([x.shape[0], 6, 4])

    for i in range(x.shape[0]):
      out[i,:,:] = torch.Tensor([
        [cos_x2[i], 0.0, 0.0, 0.0],
        [sin_x2[i], 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, cos_x5[i], 0.0],
        [0.0, 0.0, sin_x5[i], 0.0],
        [0.0, 0.0, 0.0, 1.0]])
    return out

def myFcn2(x):
    cos_x2 = torch.cos(x[:,2])
    sin_x2 = torch.sin(x[:,2])
    cos_x5 = torch.cos(x[:,5])
    sin_x5 = torch.sin(x[:,5])
    out = torch.zeros([x.shape[0], 6, 4])
    out[:, 0, 0] = cos_x2
    out[:, 1, 0] = sin_x2
    out[:, 3, 2] = cos_x5
    out[:, 4, 2] = sin_x5
    out[:, 2, 1] = 1.0
    out[:, 5, 3] = 1.0
    return out

x = torch.randn(512, 6)
t1 = time.time()
out1 = myFcn(x)
t2 = time.time()
out2 = myFcn2(x)
t3 = time.time()
print(torch.allclose(out1, out2))
print(t2-t1, t3-t2)

True
0.007179975509643555 0.00020051002502441406
1 Like