# Assigning values to Tensor without loops

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``````

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