import torch


def roll_pitch_yaw_to_rotation_matrices(roll_pitch_yaw):
  """Converts roll-pitch-yaw angles to rotation matrices.

  Args:
    roll_pitch_yaw: Tensor with shape [..., 3]. The last dimension contains
      the roll, pitch, and yaw angles in radians.  The resulting matrix
      rotates points by first applying roll around the x-axis, then pitch 
      around the y-axis, then yaw around the z-axis.

  Returns:
     Tensor with shape [..., 3, 3]. The 3x3 rotation matrices corresponding to
     the input roll-pitch-yaw angles.
  """

  cosines = torch.cos(roll_pitch_yaw)
  sines = torch.sin(roll_pitch_yaw)
  cx, cy, cz = torch.unbind(cosines, dim=-1)
  sx, sy, sz = torch.unbind(sines, dim=-1)
  # pyformat: disable
  rotation = torch.stack(
      [cz * cy, cz * sy * sx - sz * cx, cz * sy * cx + sz * sx,
       sz * cy, sz * sy * sx + cz * cx, sz * sy * cx - cz * sx,
       -sy, cy * sx, cy * cx], dim=-1)
  # pyformat: enable
  #shape = torch.cat([roll_pitch_yaw.shape[:-1], [3, 3]], axis=0)
  shape = list(roll_pitch_yaw.shape[:-1]) + [3, 3]
  rotation = torch.reshape(rotation, shape)
  return rotation