import torch


def euler_angles_to_rotation_matrices(angles):
    """
    Arguments:
    ---------
        angles: Tensor with size Kx3, where K is the number of Euler angles we
                want to transform to rotation matrices

    Returns:
    -------
        rotation_matrices: Tensor with size Kx3x3, that contains the computed
                           rotation matrices
    """
    K = angles.shape[0]
    # Allocate memory for a Tensor of size Kx3x3 that will hold the rotation
    # matrix along the x-axis
    r_x = angles.new_zeros((K, 3, 3))
    r_x[:, 0, 0] = 1.0
    c = torch.cos(angles[:, 0])
    s = torch.sin(angles[:, 0])
    r_x[torch.arange(K), 1, 1] = c
    r_x[torch.arange(K), 2, 2] = c
    r_x[torch.arange(K), 1, 2] = -s
    r_x[torch.arange(K), 2, 1] = s

    # Similar for the rotation matrices along the y-axis and z-axis
    r_y = angles.new_zeros((K, 3, 3))
    r_y[:, 1, 1] = 1.0
    c = torch.cos(angles[:, 1])
    s = torch.sin(angles[:, 1])
    r_y[torch.arange(K), 0, 0] = c
    r_y[torch.arange(K), 2, 2] = c
    r_y[torch.arange(K), 2, 0] = -s
    r_y[torch.arange(K), 0, 2] = s

    r_z = angles.new_zeros((K, 3, 3))
    r_z[:, 2, 2] = 1.0
    c = torch.cos(angles[:, 2])
    s = torch.sin(angles[:, 2])
    r_z[torch.arange(K), 0, 0] = c
    r_z[torch.arange(K), 1, 1] = c
    r_z[torch.arange(K), 0, 1] = -s
    r_z[torch.arange(K), 1, 0] = s

    return r_z.bmm(r_y.bmm(r_x))