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