TensorFlow (Python) code to add two axis-angle rotation vectors

tf_add_axis_angle_rotations.py

def add_axis_angle_rotations(rv1, rv2):
    """ Given angle*axis rotation vectors a*l and b*m, result angle*axis: c*n
    cos(c/2) = cos(a/2)cos(b/2) - sin(a/2)sin(b/2) (l . m)
    sin(c/2) (n) =
      sin(a/2)cos(b/2) (l) + cos(a/2)sin(b/2) (m) + sin(a/2)sin(b/2) (l x m)
    Args:
        rv1, rv2: Axis-angle rotation Tensors with shape [batch_size, 3]
    Returns:
        Axis-angle rotation in [batch_size, 3] Tensor which when applied is 
        equivalent to application of rv2 then rv1
    """
    # rv1, rv2 should be shaped [batch, 3] 
    a = tf.norm(rv1, axis=1, keepdims=True)
    b = tf.norm(rv2, axis=1, keepdims=True)

    l = rv1 / a
    m = rv2 / b
    cos_half_c = tf.cos(a/2) * tf.cos(b/2) - (
        tf.sin(a/2) * tf.sin(b/2) *
        tf.reduce_sum(l * m, axis=1, keepdims=True))
    sin_half_c_n = tf.sin(a/2) * tf.cos(b/2) * l + (
        tf.cos(a/2) * tf.sin(b/2) * m +
        tf.sin(a/2) * tf.sin(b/2) * tf.cross(l, m))

    half_c = tf.acos(cos_half_c)
    n = sin_half_c_n / tf.sin(half_c)
    result = n * half_c * 2

    result = tf.where(tf.equal(tf.squeeze(a), 0.0), rv2, result)
    result = tf.where(tf.equal(tf.squeeze(b), 0.0), rv1, result)
    return result