Hi again, guys!
I have some updates on this. I’ve implemented all this stuff on top of torch.script, but got stuck with a foor loop. This works (both cpu, gpu), but could be further optimized if only there is a parallel_for loop. Any ideas or plans to implement?

  
      github.com
  
  
    ferrine/geoopt/blob/aae0eb1c64a63fee0afdc2cef0592bdebf3a4492/geoopt/util/linalg.py#L20 5

# prolonged here:
if x.dim() == 2:
    # 17 milliseconds on my mac to check that condition, that is low overhead
    return torch.svd(x)
else:
    batches = x.shape[:-2]
    other = x.shape[-2:]
    flat = x.view((-1,) + other)
    slices = flat.unbind(0)
    U, D, V = [], [], []
    # I wish I had a parallel_for
    for i in range(flat.shape[0]):
        u, d, v = torch.svd(slices[i])
        U.append(u)
        D.append(d)
        V.append(v)
    U = torch.stack(U).view(batches + U[0].shape)
    D = torch.stack(D).view(batches + D[0].shape)
    V = torch.stack(V).view(batches + V[0].shape)
    return U, D, V