:param ray_o, ray_d: [..., 3] :param fg_z_max: [...,] :param fg_z_vals, bg_z_vals: [..., N_samples] :return