import torch


def rotation_matrix(angle, axis, device=None):
    """
    Return a new tensor filled with rotation matrices.

    Args:
        angle (`torch.FloatTensor`): rotation angles.
        axis (str): rotation axis, e.g. 'x' or 'y' or 'z'.
        device (`torche.device`, optional): computation device. Defaults to None.

    Returns:
        (`torch.FloatTensor`): rotation matrices.
    """
    R = torch.zeros(angle.nelement(), 3, 3, device=device)
    cos = torch.cos(angle)
    sin = torch.sin(angle)

    if axis == "x":
        R[..., 0, 0] = 1.0
        R[..., 1, 1] = cos
        R[..., 1, 2] = -sin
        R[..., 2, 1] = sin
        R[..., 2, 2] = cos

    if axis == "y":
        R[..., 0, 0] = cos
        R[..., 0, 2] = sin
        R[..., 2, 0] = -sin
        R[..., 2, 2] = cos
        R[..., 1, 1] = 1.0

    if axis == "z":
        R[..., 0, 0] = cos
        R[..., 0, 1] = -sin
        R[..., 1, 0] = sin
        R[..., 1, 1] = cos
        R[..., 2, 2] = 1.0

    return R