nan Gradient in kornia.geometry.rotation_matrix_to_quaternion

Describe the bug

When using torch.float32 the transformation from rotation matrices to quaternions outputs nan valued gradients. This is caused by trace + 1.0 (in sq = ...) which might be slightly below zero, for single precision. It should be sq = torch.sqrt(trace + 1.0 + eps) * 2.0 instead.

    def trace_positive_cond():
        sq = torch.sqrt(trace + 1.0) * 2.0  # sq = 4 * qw.
        qw = 0.25 * sq
        qx = safe_zero_division(m21 - m12, sq)
        qy = safe_zero_division(m02 - m20, sq)
        qz = safe_zero_division(m10 - m01, sq)
        if order == QuaternionCoeffOrder.XYZW:
            return torch.cat((qx, qy, qz, qw), dim=-1)
        return torch.cat((qw, qx, qy, qz), dim=-1)

Reproduction steps

import kornia
import torch


def main():
    R = torch.tensor([[1, 0, 0],
                      [0, -1, 0],
                      [0, 0, -1]], dtype=torch.float32, requires_grad=True)
    q = kornia.geometry.rotation_matrix_to_quaternion(R)
    print(torch.autograd.grad(q.sum(), R))


if __name__ == '__main__':
    main()

Expected behavior

The code above outputs a gradient including nan values. As the rotation matrix to quaternion transform includes safe divisions, it should not yield nan values.

(tensor([[    nan,  0.2500,  0.2500],
        [ 0.2500,     nan, -0.2500],
        [ 0.2500,  0.2500,     nan]]),)

Environment

python 3.8.10
pytorch 1.11.0
kornia 0.6.2

Additional context

No response

[simon-schaefer]

I have submitted the PR for adding the eps.
https://github.com/kornia/kornia/pull/1665

[ducha-aiki]

Solved in #1665