quaternion_from_rotation_matrix NaN or +Inf errors

I am trying to convert a 6D continuous representation of a rotation matrix into an axis-angle representation in a network. During backprop, I can occasionally encounter NaNs / +Inf. I was able to determine the source is quaternion_from_rotation_matrix by using a combination of tf.debugging.enable_check_numerics()and the tensorflow_graphics debugging mode.

Sometimes the tfg debugger catches an improperly normalized rotation matrix (i.e., assert_rotation_matrix_normalized), but other times it does not and quaternion_from_rotation_matrix errors in one of the safe_ops.safe_unsigned_div ops or when computing the sqrt of the trace where the trace < -1 by a very small delta.

In converting the continuous representation to a rotation matrix, I'm using an implementation very similar to this one.

Do you have any idea if this is a bug or if I might be encountering some kind of numerical instability? I am working on a repro but having some trouble freezing the model state when the error occurs to extract the bad rotation matrix.

Start with tensorflow_graphics/geometry/transformation/quaternion.py, especially the sqrt-of-trace logic around line 315 and the safe_ops.safe_unsigned_div calls. Reproduce the NaN or +Inf with the reported 6D rotation-matrix conversion and TensorFlow Graphics debugging checks. Done means determining whether the failure is numerical instability or a bug and documenting or fixing the confirmed cause.