def vector_weights(vector1: type_alias.TensorLike, vector2: type_alias.TensorLike, percent: Union[type_alias.Float, type_alias.TensorLike], eps: Optional[type_alias.Float] = None, name: str = "vector_weights") -> Tuple[tf.Tensor, tf.Tensor]: """Spherical linear interpolation (slerp) between two unnormalized vectors. This function applies geometric slerp to unnormalized vectors by first normalizing them to return the interpolation weights. It reduces to lerp when input vectors are exactly anti-parallel. Note: In the following, A1 to An are optional batch dimensions. Args: vector1: A tensor of shape `[A1, ... , An, M]`, which stores a normalized vector in its last dimension. vector2: A tensor of shape `[A1, ... , An, M]`, which stores a normalized vector in its last dimension. percent: A `float` or tensor with shape broadcastable to the shape of input vectors. eps: A small float for operation safety. If left None, its value is automatically selected using dtype of input vectors. name: A name for this op. Defaults to "vector_weights". Raises: ValueError: if the shape of `vector1`, `vector2`, or `percent` is not supported. Returns: Two tensors of shape `[A1, ... , An, 1]`, representing interpolation weights for each input vector. """ with tf.name_scope(name): vector1 = tf.convert_to_tensor(value=vector1) vector2 = tf.convert_to_tensor(value=vector2) percent = tf.convert_to_tensor(value=percent, dtype=vector1.dtype) if percent.shape.ndims == 0: percent = tf.expand_dims(percent, axis=0) shape.compare_dimensions( tensors=(vector1, vector2), axes=-1, tensor_names=("vector1", "vector2")) shape.compare_batch_dimensions( tensors=(vector1, vector2, percent), last_axes=(-2, -2, -1), broadcast_compatible=True, tensor_names=("vector1", "vector2", "percent")) normalized1 = tf.nn.l2_normalize(vector1, axis=-1) normalized2 = tf.nn.l2_normalize(vector2, axis=-1) dot_product = _safe_dot(normalized1, normalized2, eps) theta = tf.acos(dot_product) scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps) scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps) return scale1, scale2
def test_safe_sinpx_div_sinx(self): """Tests for edge cases and continuity for sin(px)/sin(x).""" angle_step = np.pi / 16.0 with self.subTest(name="all_angles"): theta = tf.range(-2.0 * np.pi, 2.0 * np.pi + angle_step / 2.0, angle_step) factor = np.random.uniform(size=(1, )) division = safe_ops.safe_sinpx_div_sinx(theta, factor) division_l = safe_ops.safe_sinpx_div_sinx(theta + 1e-10, factor) division_r = safe_ops.safe_sinpx_div_sinx(theta - 1e-10, factor) self.assertAllClose(division, division_l, rtol=1e-9) self.assertAllClose(division, division_r, rtol=1e-9) with self.subTest(name="theta_is_zero"): theta = 0.0 factor = tf.range(0.0, 1.0, 0.001) division = safe_ops.safe_sinpx_div_sinx(theta, factor) division_l = safe_ops.safe_sinpx_div_sinx(theta + 1e-10, factor) division_r = safe_ops.safe_sinpx_div_sinx(theta - 1e-10, factor) self.assertAllClose(division, division_l, atol=1e-9) self.assertAllClose(division, division_r, atol=1e-9) # According to l'Hopital rule, limit should be factor self.assertAllClose(division, factor, atol=1e-9) with self.subTest(name="theta_is_pi"): theta = np.pi factor = tf.range(0.0, 1.001, 0.001) division = safe_ops.safe_sinpx_div_sinx(theta, factor) division_l = safe_ops.safe_sinpx_div_sinx(theta + 1e-10, factor) division_r = safe_ops.safe_sinpx_div_sinx(theta - 1e-10, factor) self.assertAllClose(division, division_l, atol=1e-9) self.assertAllClose(division, division_r, atol=1e-9)
def quaternion_weights( quaternion1: type_alias.TensorLike, quaternion2: type_alias.TensorLike, percent: Union[type_alias.Float, type_alias.TensorLike], eps: Optional[type_alias.Float] = None, name: str = "quaternion_weights") -> Tuple[tf.Tensor, tf.Tensor]: """Calculates slerp weights for two normalized quaternions. Given a percent and two normalized quaternions, this function returns the slerp weights. It can also produce extrapolation weights when percent is outside of the [0, 1] range. It reduces to lerp when input quaternions are almost parallel or anti-parallel. Input quaternions are assumed to be normalized. The tf.graphics debug flag TFG_ADD_ASSERTS_TO_GRAPH defined in tfg_flags.py can be set to add assertions to the graph that check whether the inputs are normalized, and whether Inf or Nan values are produced. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion1: A tensor of shape `[A1, ... , An, 4]` storing normalized quaternions in its last dimension. quaternion2: A tensor of shape `[A1, ... , An, 4]` storing normalized quaternions in its last dimension. percent: A `float` or a tensor with a shape broadcastable to the shape `[A1, ... , An]`. eps: A `float` used to make operations safe. When left as None, the function automatically picks the best epsilon based on the dtype and the operation. name: A name for this op. Defaults to "quaternion_weights". Raises: ValueError: If the shapes of quaternions do not match, if the last dimensions of quaternions are not 4, or if percent is neither a float, nor a tensor with last dimension 1. Returns: Two tensors of shape `[A1, ... , An, 1]` each, which are the two slerp weights for each quaternion. """ with tf.name_scope(name): quaternion1 = tf.convert_to_tensor(value=quaternion1) quaternion2 = tf.convert_to_tensor(value=quaternion2) percent = tf.convert_to_tensor(value=percent, dtype=quaternion1.dtype) if percent.shape.ndims == 0: percent = tf.expand_dims(percent, axis=0) shape.check_static( tensor=quaternion1, tensor_name="quaternion1", has_dim_equals=(-1, 4)) shape.check_static( tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4)) shape.compare_batch_dimensions( tensors=(quaternion1, quaternion2, percent), last_axes=(-2, -2, -1), broadcast_compatible=True, tensor_names=("quaternion1", "quaternion2", "percent")) quaternion1 = asserts.assert_normalized(quaternion1) quaternion2 = asserts.assert_normalized(quaternion2) dot_product = _safe_dot(quaternion1, quaternion2, eps) # Take the shorter path theta = tf.acos(tf.abs(dot_product)) # safe_sinpx_div_sinx returns p for very small x, which means slerp reduces # to lerp automatically. scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps) scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps) # Flip the sign of scale1 if quaternions are in different hemispheres. # tf.sign can make scale1 zero if quaternions are orthogonal. scale1 *= safe_ops.nonzero_sign(dot_product) return scale1, scale2