Relax slerp.interpolate() input shape test and make it coherent with docstring.

PiperOrigin-RevId: 503185356

Author: tensorflower-gardener

tensorflow_graphics/math/interpolation/slerp.py

@@ -170,8 +170,8 @@ def quaternion_weights(
       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]`.
+    percent: A `float` or tensor with shape broadcastable to the shape of input
+      vectors.
     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".
@@ -198,7 +198,7 @@ def quaternion_weights(
         tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4))
     shape.compare_batch_dimensions(
         tensors=(quaternion1, quaternion2, percent),
-        last_axes=(-2, -2, -1),
+        last_axes=-1,
         broadcast_compatible=True,
         tensor_names=("quaternion1", "quaternion2", "percent"))
     quaternion1 = asserts.assert_normalized(quaternion1)
@@ -266,7 +266,7 @@ def vector_weights(vector1: type_alias.TensorLike,
         tensor_names=("vector1", "vector2"))
     shape.compare_batch_dimensions(
         tensors=(vector1, vector2, percent),
-        last_axes=(-2, -2, -1),
+        last_axes=-1,
         broadcast_compatible=True,
         tensor_names=("vector1", "vector2", "percent"))
     normalized1 = tf.nn.l2_normalize(vector1, axis=-1)

tensorflow_graphics/math/interpolation/tests/slerp_test.py

@@ -127,7 +127,7 @@ def test_unnormalized_quaternion_weights_exception_raised(self):
   @parameterized.parameters(
       ((4,), (4,), (1,)),
       ((None, 4), (None, 4), (None, 1)),
-      ((None, 4), (None, 4), (None, 4)),
+      ((5, 1, 4), (5, 1, 4), (3, 1)),
   )
   def test_quaternion_weights_exception_not_raised(self, *shapes):
     """Tests that valid input shapes do not raise exceptions for qslerp."""
@@ -140,6 +140,8 @@ def test_quaternion_weights_exception_not_raised(self, *shapes):
        (1,)),
       ("Not all batch dimensions are broadcast-compatible.", (1, 4), (3, 4),
        (2,)),
+      ("Not all batch dimensions are broadcast-compatible.", (5, 1, 4),
+       (5, 1, 4), (3,)),
   )
   def test_quaternion_weights_exception_raised(self, error_msg, *shapes):
     """Tests that the shape exceptions are properly raised for qslerp."""