import torch


def quat_distance(target_quats, pred_quats, eps=1e-6, sign_invariant=True):
    """Angle of the "difference rotation" between target_quats and pred_quats.

    Though arguments `target_quats` and `pred_quats` are suggestively named,
    the function is invariant to the actual ordering of these arguments.
    Args:
      target_quats: the ground truth quaternion orientations: shape=(..., 4).
      pred_quats: the predicted quaternion orientations: shape=(..., 4).
      eps: tolerance parameter to ensure arccos evaluates correctly.
      sign_invariant: quaternions q and -q specify equivalent orientations,
        however their geodesic distance is not 0. This argument specifies whether
        we consider the sign information or whether we pick the sign which
        gives us the shortest distance between true and predicted quaternions.

    Returns:
      shape=(..., 1). Angle subtended by true and predicted quaternions along a
      great arc of the S^3 sphere (also equivalent to double the geodesic
      distance).
      If sign_invariant=True then the minimum such angle/distance is returned by
      ignoring the sign of the quaternions.
    """
    quat_dot = torch.sum(target_quats * pred_quats, -1, keepdim=True)
    if sign_invariant:
        quat_dot = torch.abs(quat_dot) - eps
    else:
        quat_dot = quat_dot - eps * torch.sign(quat_dot)
    return 2 * torch.acos(quat_dot)