import torch


def compute_rays_length(rays_d):
    """Compute ray length.

    Args:
        rays_d: [R, 3] float tensor. Ray directions.

    Returns:
        rays_length: [R, 1] float tensor. Ray lengths.
    """
    rays_length = torch.norm(rays_d, dim=-1, keepdim=True)  # [N_rays, 1]
    return rays_length