import torch


def ray_to_box_coordinate_frame_pairwise(box_center, box_rotation_matrix,
                                         rays_start_point, rays_end_point):
    """Moves a set of rays into a box's coordinate frame.

    Args:
        box_center: A tensor of size [3] or [r, 3].
        box_rotation_matrix: A tensor of size [3, 3] or [r, 3, 3].
        rays_start_point: A tensor of size [r, 3] where r is the number of rays.
        rays_end_points: A tensor of size [r, 3] where r is the number of rays.

    Returns:
        rays_start_point_in_box_frame: A tensor of size [r, 3].
        rays_end_point_in_box_frame: A tensor if size [r, 3].
    """
    r = rays_start_point.size()[0]
    box_center = torch.broadcast_to(box_center, (r, 3))
    box_rotation_matrix = torch.broadcast_to(box_rotation_matrix, (r, 3, 3))
    rays_start_point_in_box_frame = torch.matmul(
        (rays_start_point - box_center).unsqueeze(1),
        box_rotation_matrix)
    rays_end_point_in_box_frame = torch.matmul(
        (rays_end_point - box_center).unsqueeze(1),
        box_rotation_matrix)
    return (rays_start_point_in_box_frame.view(-1, 3), 
            rays_end_point_in_box_frame.view(-1, 3))