import torch


def _build_matrix_from_sines_and_cosines(sin_angles, cos_angles):
    """Builds a rotation matrix from sines and cosines of Euler angles.

Note:
  In the following, A1 to An are optional batch dimensions.

Args:
  sin_angles: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
    represents the sine of the Euler angles.
  cos_angles: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
    represents the cosine of the Euler angles.

Returns:
  A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions
  represent a 3d rotation matrix.
"""
    # sin_angles.shape.assert_is_compatible_with(cos_angles.shape)
    sx, sy, sz = torch.unbind(sin_angles, dim=-1)
    cx, cy, cz = torch.unbind(cos_angles, dim=-1)
    m00 = cy * cz
    m01 = (sx * sy * cz) - (cx * sz)
    m02 = (cx * sy * cz) + (sx * sz)
    m10 = cy * sz
    m11 = (sx * sy * sz) + (cx * cz)
    m12 = (cx * sy * sz) - (sx * cz)
    m20 = -sy
    m21 = sx * cy
    m22 = cx * cy
    matrix = torch.stack((m00, m01, m02, m10, m11, m12, m20, m21, m22), dim=-1)  # pyformat: disable
    #output_shape = torch.cat((sin_angles.shape[:-1], (3, 3)), dim=-1)
    shape = sin_angles.shape[:-1] + (3, 3)
    return matrix.reshape(shape)