import torch


def sample_points(rays_o, rays_d, near, far, num_samples, perturb=False):
    """
    Sample points along the ray
    Args:
        rays_o (num_rays, 3): ray origins
        rays_d (num_rays, 3): ray directions
        near (float): near plane
        far (float): far plane
        num_samples (int): number of points to sample along each ray
        perturb (bool): if True, use randomized stratified sampling
    Returns:
        t_vals (num_rays, num_samples): sampled t values
        coords (num_rays, num_samples, 3): coordinate of the sampled points
    """
    num_rays = rays_o.shape[0]

    t_vals = torch.linspace(near, far, num_samples, device=rays_o.device)
    t_vals = t_vals.expand(num_rays, num_samples)   # t_vals has shape (num_samples)
                                                    # we must broadcast it to (num_rays, num_samples)
    if perturb:
        rand = torch.rand_like(t_vals) * (far-near)/num_samples
        t_vals = t_vals + rand

    coords = rays_o.unsqueeze(dim=-2) + t_vals.unsqueeze(dim=-1) * rays_d.unsqueeze(dim=-2)

    return t_vals, coords