Optimize for loops in torch.jit

Etienne_Perot · February 6, 2022, 7:50am

Here is a code that is super slow (100ms !!!):

gist.github.com

https://gist.github.com/etienne87/24c2a42d5816cfe61c1e1444f3a42a5d

double_reflection_loop_torch.py

torch.jit.script
def drloop(V: torch.Tensor,T: torch.Tensor) -> torch.Tensor:
    n = len(V)
    Ri = torch.zeros((3,), dtype=torch.float32, device=V.device)
    Ri[0] = -T[0][1]
    Ri[1] = T[0][0]
    Ri[2] = 0
    Ris = [Ri[None]]
    
    for i in range(n - 1):

This file has been truncated. show original

Basically it is a for loop with super simple dot products over vectors of 3 components.

How to optimize this? should i compute the backward myself instead?