Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Root decomposition fails for simple matrices

See original GitHub issue

Observing the following:

I = torch.eye(4)
NonLazyTensor(I).root_decomposition()

# this will often return nans, and sometimes an approx. solution
tensor([[nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan]])

Inoise = I + 1e-3 * torch.diag(torch.rand(4))
NonLazyTensor(Inoise).root_decomposition()

tensor([[ 1.6966e-05, -1.0755e-05,  1.0002e+00, -1.2987e-04],
        [-1.0000e+00, -1.4206e-03,  1.3592e-05,  5.5494e-06],
        [ 1.4198e-03, -1.0000e+00, -1.9231e-06, -1.2009e-05],
        [-5.9421e-06,  1.3196e-05, -1.0358e-04, -1.0005e+00]])

Ddiag = torch.ones(4)
DiagLazyTensor(Ddiag).root_decomposition()

# again most of the time this is nan and sometimes an approx solution
tensor([[nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan]])

Ddiagnoise = Ddiag  + 1e-3 * torch.rand(4)
DiagLazyTensor(Ddiagnoise).root_decomposition()

tensor([[ 0.0002,  1.0002,  0.0010, -0.0000],
        [ 0.9999,  0.0002, -0.0001,  0.0001],
        [ 0.0002,  0.0002, -0.0004, -1.0006],
        [-0.0003, -0.0010,  1.0003, -0.0000]])

At least for DiagLazyTensor we should just short-circuit the special-case and just return DiagLazyTensor(torch.sqrt(input_diag)). Not sure about the other cases, things seem to work reasonably if the matrices have off-diagonal elements.