jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def surround_views(initial_pose, angle_amp=1.0, num_frames=60):
    rad = torch.from_numpy(
        np.linspace(0, 2 * np.pi, num=num_frames, endpoint=False, dtype=np.float32))

    initial_pos = initial_pose[:3, -1]
    initial_pos_dist = torch.norm(initial_pos)
    initial_pos_norm = initial_pos / initial_pos_dist
    initial_angle = torch.asin(initial_pos_norm[-1])

    angles = initial_angle * (rad.sin() * angle_amp + 1)
    pos_xy = F.normalize(initial_pos_norm[:2], dim=0) @ torch.stack(
        [rad.cos(), -rad.sin(),
         rad.sin(), rad.cos()], dim=-1).reshape(-1, 2, 2)
    pos = torch.cat(
        [pos_xy * angles.cos().unsqueeze(-1), angles.sin().unsqueeze(-1)],
        dim=-1) * initial_pos_dist
    rot = look_at(pos, torch.zeros_like(pos), pos.new_tensor([0, 0, 1]).expand(pos.size()))
    poses = torch.cat(
        [torch.cat([rot, pos.unsqueeze(-1)], dim=-1),
         rot.new_tensor([0, 0, 0, 1]).expand(num_frames, 1, -1)], dim=-2)

    return poses
//  [python]
def surround_views(initial_pose, angle_amp=1.0, num_frames=60):
    rad = torch.from_numpy(
        np.linspace(0, 2 * np.pi, num=num_frames, endpoint=False, dtype=np.float32))

    initial_pos = initial_pose[:3, -1]
    initial_pos_dist = torch.linalg.norm(initial_pos)
    initial_pos_norm = initial_pos / initial_pos_dist
    initial_angle = torch.asin(initial_pos_norm[-1])

    angles = initial_angle * (rad.sin() * angle_amp + 1)
    pos_xy = F.normalize(initial_pos_norm[:2], dim=0) @ torch.stack(
        [rad.cos(), -rad.sin(),
         rad.sin(), rad.cos()], dim=-1).reshape(-1, 2, 2)
    pos = torch.cat(
        [pos_xy * angles.cos().unsqueeze(-1), angles.sin().unsqueeze(-1)],
        dim=-1) * initial_pos_dist
    rot = look_at(pos, torch.zeros_like(pos), pos.new_tensor([0, 0, 1]).expand(pos.size()))
    poses = torch.cat(
        [torch.cat([rot, pos.unsqueeze(-1)], dim=-1),
         rot.new_tensor([0, 0, 0, 1]).expand(num_frames, 1, -1)], dim=-2)

    return poses
//  [python]
def normalise(x: Tensor, p: int = 2) -> Tensor:
    """
    Args:
        x: A 2D tensor
        p: Specifies the exact p-norm

    Returns:
        Normalised input

    """
    assert x.ndim == 2
    xn = torch.norm(x, p, dim=1).detach()
    x = x.div(xn.unsqueeze(1))
    return x
//  [python]
def normalise(x: Tensor, p: int = 2) -> Tensor:
    """
    Args:
        x: A 2D tensor
        p: Specifies the exact p-norm

    Returns:
        Normalised input

    """
    assert x.ndim == 2
    xn = torch.linalg.norm(x, p, dim=1).detach()
    x = x.div(xn.unsqueeze(1))
    return x
//  [python]
def cos_between_vectors(x, y, do_clamp=True):
    """
    Args:
        x: (..., 3)
        y: (..., 3)
    Return:
        (...), cos theta between two vectors
    """
    cos_theta = torch.sum(x * y, dim=-1) / (
        torch.norm(x, dim=-1) * torch.norm(y, dim=-1)
    )
    if do_clamp == True:
        return torch.clamp(cos_theta, -1, 1)
    else:
        return cos_theta
//  [python]
def cos_between_vectors(x, y, do_clamp=True):
    """
    Args:
        x: (..., 3)
        y: (..., 3)
    Return:
        (...), cos theta between two vectors
    """
    cos_theta = torch.sum(x * y, dim=-1) / (
        torch.linalg.norm(x, dim=-1) * torch.linalg.norm(y, dim=-1)
    )
    if do_clamp == True:
        return torch.clamp(cos_theta, -1, 1)
    else:
        return cos_theta
//  [python]
def dihedral_from_four_points(p0, p1, p2, p3):
    """
    Args:
        p0-3:   (*, 3).
    Returns:
        Dihedral angles in radian, (*, ).
    """
    v0 = p2 - p1
    v1 = p0 - p1
    v2 = p3 - p2
    u1 = torch.cross(v0, v1, dim=-1)
    n1 = u1 / torch.norm(u1, dim=-1, keepdim=True)
    u2 = torch.cross(v0, v2, dim=-1)
    n2 = u2 / torch.norm(u2, dim=-1, keepdim=True)
    sgn = torch.sign((torch.cross(v1, v2, dim=-1) * v0).sum(-1))
    dihed = sgn * torch.acos((n1 * n2).sum(-1).clamp(min=-0.999999, max=0.999999))
    dihed = torch.nan_to_num(dihed)
    return dihed
//  [python]
def dihedral_from_four_points(p0, p1, p2, p3):
    """
    Args:
        p0-3:   (*, 3).
    Returns:
        Dihedral angles in radian, (*, ).
    """
    v0 = p2 - p1
    v1 = p0 - p1
    v2 = p3 - p2
    u1 = torch.cross(v0, v1, dim=-1)
    n1 = u1 / torch.linalg.norm(u1, dim=-1, keepdim=True)
    u2 = torch.cross(v0, v2, dim=-1)
    n2 = u2 / torch.linalg.norm(u2, dim=-1, keepdim=True)
    sgn = torch.sign((torch.cross(v1, v2, dim=-1) * v0).sum(-1))
    dihed = sgn * torch.acos((n1 * n2).sum(-1).clamp(min=-0.999999, max=0.999999))
    dihed = torch.nan_to_num(dihed)
    return dihed
//  [python]
def plane_pts(rays, planes, id_planes, near, method="planes"):
    """Ray-Plane intersection for given planes in the scene

    Args:
        rays: ray origin and directions
        planes: first plane position, plane normal and distance between planes
        id_planes: ids of used planes
        near: distance between camera pose and first intersecting plane
        method: Method used

    Returns:
        pts: [N_rays, N_samples+N_importance] - intersection points of rays and selected planes
        z_vals: position of the point along each ray respectively
    """
    # Extract ray and plane definitions
    rays_o, rays_d = rays
    N_rays = rays_o.shape[0]
    # N_rays = rays_o.get_shape().as_list()[0]
    plane_bds, plane_normal, delta = planes
    plane_bds = torch.from_numpy(plane_bds).to(torch.float32)
    plane_normal = torch.from_numpy(plane_normal).to(torch.float32)

    # Get amount of all planes
    n_planes = torch.ceil(torch.norm(plane_bds[:, -1] - plane_bds[:, 0]) / delta) + 1
    # n_planes = tf.math.ceil(tf.norm(plane_bds[:, -1] - plane_bds[:, 0]) / delta) + 1

    # Calculate how far the ray_origins lies apart from each plane
    d_ray_first_plane = torch.matmul(plane_bds[:, 0] - rays_o, plane_normal[:, None])
    d_ray_first_plane = torch.maximum(-d_ray_first_plane, -near)
    # d_ray_first_plane = tf.matmul(plane_bds[:, 0]-rays_o, plane_normal[:, None])
    # d_ray_first_plane = tf.maximum(-d_ray_first_plane, -near)

    # Get the ids of the planes in front of each ray starting from near distance upto the far plane
    start_id = torch.ceil((d_ray_first_plane + near) / delta)
    plane_id = start_id + torch.from_numpy(id_planes).to(torch.float32)
    if method == "planes":
        plane_id = torch.cat([plane_id[:, :-1], n_planes.repeat_interleave(N_rays)[:, None]], dim=1)
        # plane_id = tf.concat([plane_id[:, :-1], tf.repeat(n_planes, N_rays)[:, None]], axis=1)
    elif method == "planes_plus":
        # Experimental setup, that got discarded due to lower or the same quality
        plane_id = torch.cat(
            [
                plane_id[:, :1],
                id_planes[None, 1:-1].repeat_interleave(N_rays, dim=0),
                n_planes.repeat_interleave(N_rays)[:, None],
            ],
            dim=1,
        )
        # plane_id = tf.concat([plane_id[:, :1],
        #                       tf.repeat(id_planes[None, 1:-1], N_rays, axis=0),
        #                       tf.repeat(n_planes, N_rays)[:, None]], axis=1)

    # [N_samples, N_rays, xyz]
    z_planes = plane_normal[None, None, :] * torch.transpose(plane_id * delta, 0, 1)[..., None]
    # z_planes = plane_normal[None, None, :] * tf.transpose(plane_id*delta)[..., None]
    relevant_plane_origins = plane_bds[:, 0][None, None, :] + z_planes

    # Distance between each ray's origin and associated planes
    d_plane_pose = relevant_plane_origins - rays_o[None, :, :]

    n = torch.matmul(d_plane_pose, plane_normal[..., None])
    z = torch.matmul(rays_d, plane_normal[..., None])

    z_vals = torch.transpose(torch.squeeze(n / z), 0, 1)

    pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals[..., None]

    return pts, z_vals
//  [python]
def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
        if len(self.train_dataloader) == 0:
            raise ValueError("Training dataloader must not be empty.")

        if self.use_block_cg:
            return self._solve_pbcg(rhs)

        hvp = create_hvp_function(
            self.model,
            self.loss,
            self.train_dataloader,
            precompute_grad=self.precompute_grad,
        )

        def reg_hvp(v: torch.Tensor):
            return hvp(v) + self.hessian_regularization * v.type(rhs.dtype)

        y_norm = torch.norm(rhs, dim=0)

        stopping_val = torch.clamp(self.rtol**2 * y_norm, min=self.atol**2)

        batch_cg = torch.zeros_like(rhs)

        for idx, (bi, _tol) in enumerate(
            tqdm(
                zip(rhs, stopping_val),
                disable=not self.progress,
                desc="Conjugate gradient",
            )
        ):
            batch_result = self._solve_pcg(
                reg_hvp,
                bi,
                tol=_tol,
                x0=self.x0,
                maxiter=self.maxiter,
            )
            batch_cg[idx] = batch_result

        return batch_cg
//  [python]
def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
        if len(self.train_dataloader) == 0:
            raise ValueError("Training dataloader must not be empty.")

        if self.use_block_cg:
            return self._solve_pbcg(rhs)

        hvp = create_hvp_function(
            self.model,
            self.loss,
            self.train_dataloader,
            precompute_grad=self.precompute_grad,
        )

        def reg_hvp(v: torch.Tensor):
            return hvp(v) + self.hessian_regularization * v.type(rhs.dtype)

        y_norm = torch.linalg.norm(rhs, dim=0)

        stopping_val = torch.clamp(self.rtol**2 * y_norm, min=self.atol**2)

        batch_cg = torch.zeros_like(rhs)

        for idx, (bi, _tol) in enumerate(
            tqdm(
                zip(rhs, stopping_val),
                disable=not self.progress,
                desc="Conjugate gradient",
            )
        ):
            batch_result = self._solve_pcg(
                reg_hvp,
                bi,
                tol=_tol,
                x0=self.x0,
                maxiter=self.maxiter,
            )
            batch_cg[idx] = batch_result

        return batch_cg
//  [python]
def up_sample_mid(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, gradient):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, p=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]

        cos_val = (rays_d[:, None, :].expand(pts.shape) * gradient.view(pts.shape)).sum(-1)[:,:-1]

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        cdf = torch.sigmoid(prev_sdf * inv_s)
        e = inv_s * (1 - cdf) * (-cos_val) * dist
        alpha = (1 - torch.exp(-e)).reshape(batch_size, n_samples-1).clip(0.0, 1.0)

        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def up_sample_mid(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, gradient):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]

        cos_val = (rays_d[:, None, :].expand(pts.shape) * gradient.view(pts.shape)).sum(-1)[:,:-1]

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        cdf = torch.sigmoid(prev_sdf * inv_s)
        e = inv_s * (1 - cdf) * (-cos_val) * dist
        alpha = (1 - torch.exp(-e)).reshape(batch_size, n_samples-1).clip(0.0, 1.0)

        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def forward(self, flows: torch.Tensor, boxes: torch.Tensor) -> torch.Tensor:
        """Extract velocioty features by filling a histogram.

        Args:
            flows (torch.Tensor): Batch of optical flow images of shape (N, 2, H, W)
            boxes (torch.Tensor): Bounding box coordinates of shaspe (M, 5).
                First column indicates batch index of the bbox.

        Returns:
            Tensor: Velocity feature tensor of shape (M, n_bins)
        """
        flow_regions = roi_align(flows, boxes, output_size=[224, 224])

        # cartesian to polar
        mag_batch = torch.sqrt((flow_regions**2).sum(dim=1))
        theta_batch = torch.atan2(flow_regions[:, 0, ...], flow_regions[:, 1, ...])

        # compute velocity histogram
        velocity_histograms = []
        for mag, theta in zip(mag_batch, theta_batch, strict=True):
            histogram_mag = torch.histogram(
                input=theta.cpu(),
                bins=self.n_bins,
                range=(-torch.pi, torch.pi),
                weight=mag.cpu(),
            ).hist
            histogram_counts = torch.histogram(input=theta.cpu(), bins=self.n_bins, range=(-torch.pi, torch.pi)).hist
            final_histogram = torch.zeros_like(histogram_mag)
            mask = histogram_counts != 0
            final_histogram[mask] = histogram_mag[mask] / histogram_counts[mask]
            velocity_histograms.append(final_histogram)

        if len(velocity_histograms) == 0:
            return torch.empty(0, self.n_bins).to(flows.device)
        return torch.stack(velocity_histograms).to(flows.device)
//  [python]
def forward(self, flows: torch.Tensor, boxes: torch.Tensor) -> torch.Tensor:
        """Extract velocioty features by filling a histogram.

        Args:
            flows (torch.Tensor): Batch of optical flow images of shape (N, 2, H, W)
            boxes (torch.Tensor): Bounding box coordinates of shaspe (M, 5).
                First column indicates batch index of the bbox.

        Returns:
            Tensor: Velocity feature tensor of shape (M, n_bins)
        """
        flow_regions = roi_align(flows, boxes, output_size=[224, 224])

        # cartesian to polar
        mag_batch = torch.linalg.norm(flow_regions, axis=1, ord=2)
        theta_batch = torch.atan2(flow_regions[:, 0, ...], flow_regions[:, 1, ...])

        # compute velocity histogram
        velocity_histograms = []
        for mag, theta in zip(mag_batch, theta_batch, strict=True):
            histogram_mag = torch.histogram(
                input=theta.cpu(),
                bins=self.n_bins,
                range=(-torch.pi, torch.pi),
                weight=mag.cpu(),
            ).hist
            histogram_counts = torch.histogram(input=theta.cpu(), bins=self.n_bins, range=(-torch.pi, torch.pi)).hist
            final_histogram = torch.zeros_like(histogram_mag)
            mask = histogram_counts != 0
            final_histogram[mask] = histogram_mag[mask] / histogram_counts[mask]
            velocity_histograms.append(final_histogram)

        if len(velocity_histograms) == 0:
            return torch.empty(0, self.n_bins).to(flows.device)
        return torch.stack(velocity_histograms).to(flows.device)
//  [python]
def xyxy_to_pts_along_vanishing_line(args, bbox_2d, calibs, images):
    pos = xyxy_to_cxcy(bbox_2d, mode=args.vanishing_line_mode)
    # vanishing point
    v_x, v_y = calibs[:, 0, 2], calibs[:, 1, 2]
    v_pt = torch.stack([v_x, v_y]).T
    # center vanishing line
    i_center_bottom_pt = torch.stack(
        [
            torch.ones_like(v_pt[:, 0]) * images.shape[-1] * 0.5,
            # get mid width of image
            torch.zeros_like(v_pt[:, 0])
        ]
    ).T
    v_cl = v_pt - i_center_bottom_pt
    # standardize to [0,1], origin bottom left of image
    y_max = images.shape[-2]
    x_max = images.shape[-1]
    scaler = torch.tensor([x_max, y_max], device=v_cl.device)
    pos = [p / scaler.unsqueeze(0) for p in pos]
    v_pt = v_pt / scaler.unsqueeze(0)
    v_cl = v_cl / scaler.unsqueeze(0)
    # object vanishing line (from object to optical center)
    obj_v_l = [v.unsqueeze(0) - p for v, p in zip(v_pt, pos)]
    # project object vanishing line onto center vanishing line
    proj_obj_v_l = [
        o_v_l.matmul(v_l.unsqueeze(-1)) / torch.norm(v_l) for o_v_l, v_l in
        zip(obj_v_l, v_cl)
    ]
    # subtract from 0 so objects closer to the vanishing point are larger
    proj_obj_v_l = [0 - p for p in proj_obj_v_l]
    pos = [torch.stack([p[:, 0], proj.squeeze(-1)]).T for p, proj in
           zip(pos, proj_obj_v_l)]
    return pos
//  [python]
def xyxy_to_pts_along_vanishing_line(args, bbox_2d, calibs, images):
    pos = xyxy_to_cxcy(bbox_2d, mode=args.vanishing_line_mode)
    # vanishing point
    v_x, v_y = calibs[:, 0, 2], calibs[:, 1, 2]
    v_pt = torch.stack([v_x, v_y]).T
    # center vanishing line
    i_center_bottom_pt = torch.stack(
        [
            torch.ones_like(v_pt[:, 0]) * images.shape[-1] * 0.5,
            # get mid width of image
            torch.zeros_like(v_pt[:, 0])
        ]
    ).T
    v_cl = v_pt - i_center_bottom_pt
    # standardize to [0,1], origin bottom left of image
    y_max = images.shape[-2]
    x_max = images.shape[-1]
    scaler = torch.tensor([x_max, y_max], device=v_cl.device)
    pos = [p / scaler.unsqueeze(0) for p in pos]
    v_pt = v_pt / scaler.unsqueeze(0)
    v_cl = v_cl / scaler.unsqueeze(0)
    # object vanishing line (from object to optical center)
    obj_v_l = [v.unsqueeze(0) - p for v, p in zip(v_pt, pos)]
    # project object vanishing line onto center vanishing line
    proj_obj_v_l = [
        o_v_l.matmul(v_l.unsqueeze(-1)) / torch.linalg.norm(v_l) for o_v_l, v_l in
        zip(obj_v_l, v_cl)
    ]
    # subtract from 0 so objects closer to the vanishing point are larger
    proj_obj_v_l = [0 - p for p in proj_obj_v_l]
    pos = [torch.stack([p[:, 0], proj.squeeze(-1)]).T for p, proj in
           zip(pos, proj_obj_v_l)]
    return pos
//  [python]
def _get_weight_norm(self, weight, lora_weight, scaling) -> torch.Tensor:
        # calculate L2 norm of weight matrix, column-wise
        weight = transpose(weight, self.fan_in_fan_out)
        weight = weight + scaling * lora_weight
        weight_norm = torch.norm(weight, dim=1, keepdim=True).to(weight.dtype)
        return weight_norm
//  [python]
def _get_weight_norm(self, weight, lora_weight, scaling) -> torch.Tensor:
        # calculate L2 norm of weight matrix, column-wise
        weight = transpose(weight, self.fan_in_fan_out)
        weight = weight + scaling * lora_weight
        weight_norm = torch.linalg.norm(weight, dim=1).to(weight.dtype)
        return weight_norm
//  [python]
def _compute_likelihood(self):
        """
        Compute (negative log) likelihood in the linear Gaussian case.

        Return
        ------
        torch.Tensor: Likelihood term (scalar-valued).
        """
        if self.equal_variances:  # Assuming equal noise variances
            return (0.5 * self.d
                    * torch.log(torch.square(torch.norm(self.X - self.X @ self.B)))
                    - torch.linalg.slogdet(torch.eye(self.d).to(self.device) - self.B)[1])
        else:  # Assuming non-equal noise variances
            return (0.5
                    * torch.sum(torch.log(torch.sum(torch.square(self.X - self.X @ self.B), dim=0)))
                    - torch.linalg.slogdet(torch.eye(self.d).to(self.device) - self.B)[1])
//  [python]
def _compute_likelihood(self):
        """
        Compute (negative log) likelihood in the linear Gaussian case.

        Return
        ------
        torch.Tensor: Likelihood term (scalar-valued).
        """
        if self.equal_variances:  # Assuming equal noise variances
            return (0.5 * self.d
                    * torch.log(torch.square(torch.linalg.norm(self.X - self.X @ self.B)))
                    - torch.linalg.slogdet(torch.eye(self.d).to(self.device) - self.B)[1])
        else:  # Assuming non-equal noise variances
            return (0.5
                    * torch.sum(torch.log(torch.sum(torch.square(self.X - self.X @ self.B), dim=0)))
                    - torch.linalg.slogdet(torch.eye(self.d).to(self.device) - self.B)[1])
//  [python]
def torch_score(
            self,
            vec: torch.tensor,
            t: torch.tensor,
            eps: float=1e-6,
        ):
        """Computes the score of IGSO(3) density as a rotation vector.

        Same as score function but uses pytorch and performs a look-up.

        Args:
            vec: [..., 3] array of axis-angle rotation vectors.
            t: continuous time in [0, 1].

        Returns:
            [..., 3] score vector in the direction of the sampled vector with
            magnitude given by _score_norms.
        """
        omega = torch.norm(vec, dim=-1) + eps
        if self.use_cached_score:
            score_norms_t = self._score_norms[self.t_to_idx(du.move_to_np(t))]
            score_norms_t = torch.tensor(score_norms_t).to(vec.device)
            omega_idx = torch.bucketize(
                omega, torch.tensor(self.discrete_omega[:-1]).to(vec.device))
            omega_scores_t = torch.gather(
                score_norms_t, 1, omega_idx)
        else:
            sigma = self.discrete_sigma[self.t_to_idx(du.move_to_np(t))]
            sigma = torch.tensor(sigma).to(vec.device)
            omega_vals = igso3_expansion(omega, sigma[:, None], use_torch=True)
            omega_scores_t = score(omega_vals, omega, sigma[:, None], use_torch=True)
        return omega_scores_t[..., None] * vec / (omega[..., None] + eps)
//  [python]
def torch_score(
            self,
            vec: torch.tensor,
            t: torch.tensor,
            eps: float=1e-6,
        ):
        """Computes the score of IGSO(3) density as a rotation vector.

        Same as score function but uses pytorch and performs a look-up.

        Args:
            vec: [..., 3] array of axis-angle rotation vectors.
            t: continuous time in [0, 1].

        Returns:
            [..., 3] score vector in the direction of the sampled vector with
            magnitude given by _score_norms.
        """
        omega = torch.linalg.norm(vec, dim=-1) + eps
        if self.use_cached_score:
            score_norms_t = self._score_norms[self.t_to_idx(du.move_to_np(t))]
            score_norms_t = torch.tensor(score_norms_t).to(vec.device)
            omega_idx = torch.bucketize(
                omega, torch.tensor(self.discrete_omega[:-1]).to(vec.device))
            omega_scores_t = torch.gather(
                score_norms_t, 1, omega_idx)
        else:
            sigma = self.discrete_sigma[self.t_to_idx(du.move_to_np(t))]
            sigma = torch.tensor(sigma).to(vec.device)
            omega_vals = igso3_expansion(omega, sigma[:, None], use_torch=True)
            omega_scores_t = score(omega_vals, omega, sigma[:, None], use_torch=True)
        return omega_scores_t[..., None] * vec / (omega[..., None] + eps)
//  [python]
def up_sample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
    """
    Up sampling give a fixed inv_s
    copied from neus
    """
    batch_size, n_samples = z_vals.shape
    pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
    radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
    inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
    sdf = sdf.reshape(batch_size, n_samples)
    prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
    prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
    mid_sdf = (prev_sdf + next_sdf) * 0.5
    cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

    # ----------------------------------------------------------------------------------------------------------
    # Use min value of [ cos, prev_cos ]
    # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
    # robust when meeting situations like below:
    #
    # SDF
    # ^
    # |\          -----x----...
    # | \        /
    # |  x      x
    # |---\----/-------------> 0 level
    # |    \  /
    # |     \/
    # |
    # ----------------------------------------------------------------------------------------------------------
    prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
    cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
    cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
    cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

    dist = (next_z_vals - prev_z_vals)
    prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
    next_esti_sdf = mid_sdf + cos_val * dist * 0.5
    prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
    next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
    alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
    weights = alpha * torch.cumprod(
        torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
    z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
    return z_samples
//  [python]
def up_sample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
    """
    Up sampling give a fixed inv_s
    copied from neus
    """
    batch_size, n_samples = z_vals.shape
    pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
    radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
    inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
    sdf = sdf.reshape(batch_size, n_samples)
    prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
    prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
    mid_sdf = (prev_sdf + next_sdf) * 0.5
    cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

    # ----------------------------------------------------------------------------------------------------------
    # Use min value of [ cos, prev_cos ]
    # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
    # robust when meeting situations like below:
    #
    # SDF
    # ^
    # |\          -----x----...
    # | \        /
    # |  x      x
    # |---\----/-------------> 0 level
    # |    \  /
    # |     \/
    # |
    # ----------------------------------------------------------------------------------------------------------
    prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
    cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
    cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
    cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

    dist = (next_z_vals - prev_z_vals)
    prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
    next_esti_sdf = mid_sdf + cos_val * dist * 0.5
    prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
    next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
    alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
    weights = alpha * torch.cumprod(
        torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
    z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
    return z_samples
//  [python]
def add_agent(
        self,
        agent_type: AgentType,
        agent_state: StateTensor,
        agent_extent: Tensor,
        **kwargs,
    ) -> Tuple[GlyphRenderer, GlyphRenderer]:
        """Draws an agent at the given location, heading, and dimensions.

        Args:
            agent_type (AgentType): _description_
            agent_state (StateTensor): _description_
            agent_extent (Tensor): _description_
        """
        if torch.any(torch.isnan(agent_extent)):
            raise ValueError("Agent extents cannot be NaN!")

        length = agent_extent[0].item()
        width = agent_extent[1].item()

        x, y = agent_state.position.cpu().numpy()
        heading = agent_state.heading.cpu().numpy()

        agent_rect_coords, dir_patch_coords = vis_utils.compute_agent_rect_coords(
            agent_type, heading, length, width
        )

        source = {
            "x": agent_rect_coords[:, 0] + x,
            "y": agent_rect_coords[:, 1] + y,
            "type": [vis_utils.pretty_print_agent_type(agent_type)],
            "speed": [t.norm(agent_state.velocity).item()],
        }

        r = self.raw_figure.patch(
            x="x",
            y="y",
            source=source,
            **kwargs,
        )
        p = self.raw_figure.patch(
            x=dir_patch_coords[:, 0] + x, y=dir_patch_coords[:, 1] + y, **kwargs
        )

        return r, p
//  [python]
def add_agent(
        self,
        agent_type: AgentType,
        agent_state: StateTensor,
        agent_extent: Tensor,
        **kwargs,
    ) -> Tuple[GlyphRenderer, GlyphRenderer]:
        """Draws an agent at the given location, heading, and dimensions.

        Args:
            agent_type (AgentType): _description_
            agent_state (Tensor): _description_
            agent_extent (Tensor): _description_
        """
        if torch.any(torch.isnan(agent_extent)):
            raise ValueError("Agent extents cannot be NaN!")

        length = agent_extent[0].item()
        width = agent_extent[1].item()

        x, y = agent_state.position.cpu().numpy()
        heading = agent_state.heading.cpu().numpy()

        agent_rect_coords, dir_patch_coords = vis_utils.compute_agent_rect_coords(
            agent_type, heading, length, width
        )

        source = {
            "x": agent_rect_coords[:, 0] + x,
            "y": agent_rect_coords[:, 1] + y,
            "type": [vis_utils.pretty_print_agent_type(agent_type)],
            "speed": [torch.linalg.norm(agent_state.velocity).item()],
        }

        r = self.raw_figure.patch(
            x="x",
            y="y",
            source=source,
            **kwargs,
        )
        p = self.raw_figure.patch(
            x=dir_patch_coords[:, 0] + x, y=dir_patch_coords[:, 1] + y, **kwargs
        )

        return r, p
//  [python]
def _step(
        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
    ) -> TensorDictBase:
        if "goal_embedding" not in tensordict.keys():
            tensordict = self._embed_goal(tensordict)
        last_embedding_key = self.out_keys[0]
        last_embedding = tensordict.get(last_embedding_key, None)
        next_tensordict = super()._step(tensordict, next_tensordict)
        cur_embedding = next_tensordict.get(self.out_keys[0])
        if last_embedding is not None:
            goal_embedding = tensordict["goal_embedding"]
            reward = -torch.norm(cur_embedding - goal_embedding, dim=-1) - (
                -torch.norm(last_embedding - goal_embedding, dim=-1)
            )
            next_tensordict.set("reward", reward)
        return next_tensordict
//  [python]
def _step(
        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
    ) -> TensorDictBase:
        if "goal_embedding" not in tensordict.keys():
            tensordict = self._embed_goal(tensordict)
        last_embedding_key = self.out_keys[0]
        last_embedding = tensordict.get(last_embedding_key, None)
        next_tensordict = super()._step(tensordict, next_tensordict)
        cur_embedding = next_tensordict.get(self.out_keys[0])
        if last_embedding is not None:
            goal_embedding = tensordict["goal_embedding"]
            reward = -torch.linalg.norm(cur_embedding - goal_embedding, dim=-1) - (
                -torch.linalg.norm(last_embedding - goal_embedding, dim=-1)
            )
            next_tensordict.set("reward", reward)
        return next_tensordict
//  [python]
def mask_spheres_image(
    image: torch.Tensor,
    link_spheres_tensor: torch.Tensor,
    points: torch.Tensor,
    distance_threshold: float,
) -> Tuple[torch.Tensor, torch.Tensor]:

    if link_spheres_tensor.shape[0] != 1:
        assert link_spheres_tensor.shape[0] == points.shape[0]
    if len(points.shape) == 2:
        points = points.unsqueeze(0)

    robot_spheres = link_spheres_tensor.view(link_spheres_tensor.shape[0], -1, 4).contiguous()
    robot_spheres = robot_spheres.unsqueeze(-3)

    robot_radius = robot_spheres[..., 3]
    points = points.unsqueeze(-2)
    sph_distance = -1 * (
        torch.norm(points - robot_spheres[..., :3], dim=-1) - robot_radius
    )  # b, n_spheres
    distance = torch.max(sph_distance, dim=-1)[0]

    distance = distance.view(
        image.shape[0],
        image.shape[1],
        image.shape[2],
    )
    mask = torch.logical_and((image > 0.0), (distance > -distance_threshold))
    filtered_image = torch.where(mask, 0, image)
    return mask, filtered_image
//  [python]
def mask_spheres_image(
    image: torch.Tensor,
    link_spheres_tensor: torch.Tensor,
    points: torch.Tensor,
    distance_threshold: float,
) -> Tuple[torch.Tensor, torch.Tensor]:

    if link_spheres_tensor.shape[0] != 1:
        assert link_spheres_tensor.shape[0] == points.shape[0]
    if len(points.shape) == 2:
        points = points.unsqueeze(0)

    robot_spheres = link_spheres_tensor.view(link_spheres_tensor.shape[0], -1, 4).contiguous()
    robot_spheres = robot_spheres.unsqueeze(-3)

    robot_radius = robot_spheres[..., 3]
    points = points.unsqueeze(-2)
    sph_distance = -1 * (
        torch.linalg.norm(points - robot_spheres[..., :3], dim=-1) - robot_radius
    )  # b, n_spheres
    distance = torch.max(sph_distance, dim=-1)[0]

    distance = distance.view(
        image.shape[0],
        image.shape[1],
        image.shape[2],
    )
    mask = torch.logical_and((image > 0.0), (distance > -distance_threshold))
    filtered_image = torch.where(mask, 0, image)
    return mask, filtered_image
//  [python]
def upsample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def upsample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def forward(self, obs: torch.Tensor, action: torch.Tensor, next_obs: torch.Tensor, future_obs: torch.Tensor):
        goals = future_obs
        rewards = (torch.norm(obs - goals, dim=-1) < 1e-6).float()
        masks = 1.0 - rewards
        rewards = rewards - 1.0

        next_v1, next_v2 = self.value(next_obs, goals, is_target=True)
        next_v = torch.minimum(next_v1, next_v2)
        q = rewards + self.cfg.hilp_discount * masks * next_v

        v1_t, v2_t = self.value(obs, goals, is_target=True)
        v_t = (v1_t + v2_t) / 2
        adv = q - v_t

        q1 = rewards + self.cfg.hilp_discount * masks * next_v1
        q2 = rewards + self.cfg.hilp_discount * masks * next_v2
        v1, v2 = self.value(obs, goals, is_target=False)
        v = (v1 + v2) / 2

        value_loss1 = self.expectile_loss(adv, q1 - v1, self.cfg.hilp_expectile).mean()
        value_loss2 = self.expectile_loss(adv, q2 - v2, self.cfg.hilp_expectile).mean()
        value_loss = value_loss1 + value_loss2

        utils.soft_update_params(self.phi1, self.target_phi1, 0.005)
        utils.soft_update_params(self.phi2, self.target_phi2, 0.005)

        with torch.no_grad():
            phi1 = self.phi1(obs)
            self.running_mean = 0.995 * self.running_mean + 0.005 * phi1.mean(dim=0)
            self.running_std = 0.995 * self.running_std + 0.005 * phi1.std(dim=0)

        return value_loss, {
            'hilp/value_loss': value_loss,
            'hilp/v_mean': v.mean(),
            'hilp/v_max': v.max(),
            'hilp/v_min': v.min(),
            'hilp/abs_adv_mean': torch.abs(adv).mean(),
            'hilp/adv_mean': adv.mean(),
            'hilp/adv_max': adv.max(),
            'hilp/adv_min': adv.min(),
            'hilp/accept_prob': (adv >= 0).float().mean(),
        }
//  [python]
def forward(self, obs: torch.Tensor, action: torch.Tensor, next_obs: torch.Tensor, future_obs: torch.Tensor):
        goals = future_obs
        rewards = (torch.linalg.norm(obs - goals, dim=-1) < 1e-6).float()
        masks = 1.0 - rewards
        rewards = rewards - 1.0

        next_v1, next_v2 = self.value(next_obs, goals, is_target=True)
        next_v = torch.minimum(next_v1, next_v2)
        q = rewards + self.cfg.hilp_discount * masks * next_v

        v1_t, v2_t = self.value(obs, goals, is_target=True)
        v_t = (v1_t + v2_t) / 2
        adv = q - v_t

        q1 = rewards + self.cfg.hilp_discount * masks * next_v1
        q2 = rewards + self.cfg.hilp_discount * masks * next_v2
        v1, v2 = self.value(obs, goals, is_target=False)
        v = (v1 + v2) / 2

        value_loss1 = self.expectile_loss(adv, q1 - v1, self.cfg.hilp_expectile).mean()
        value_loss2 = self.expectile_loss(adv, q2 - v2, self.cfg.hilp_expectile).mean()
        value_loss = value_loss1 + value_loss2

        utils.soft_update_params(self.phi1, self.target_phi1, 0.005)
        utils.soft_update_params(self.phi2, self.target_phi2, 0.005)

        with torch.no_grad():
            phi1 = self.phi1(obs)
            self.running_mean = 0.995 * self.running_mean + 0.005 * phi1.mean(dim=0)
            self.running_std = 0.995 * self.running_std + 0.005 * phi1.std(dim=0)

        return value_loss, {
            'hilp/value_loss': value_loss,
            'hilp/v_mean': v.mean(),
            'hilp/v_max': v.max(),
            'hilp/v_min': v.min(),
            'hilp/abs_adv_mean': torch.abs(adv).mean(),
            'hilp/adv_mean': adv.mean(),
            'hilp/adv_max': adv.max(),
            'hilp/adv_min': adv.min(),
            'hilp/accept_prob': (adv >= 0).float().mean(),
        }
//  [python]
def _pairwise_cosine(
    x: torch.Tensor, y: torch.Tensor, y2: torch.Tensor
) -> torch.Tensor:
    x2 = torch.norm(x, dim=1).reshape((-1, 1))
    return 1 - (x @ y.T).div_(x2).div_(y2)
//  [python]
def _pairwise_cosine(
    x: torch.Tensor, y: torch.Tensor, y2: torch.Tensor
) -> torch.Tensor:
    x2 = torch.linalg.norm(x, dim=1).reshape((-1, 1))
    return 1 - (x @ y.T).div_(x2).div_(y2)
//  [python]
def test_prox_l21():
    prox_val = .75

    Dx = torch.randn(1, 1, 1, 1)
    Dy = torch.randn(1, 1, 1, 1)
    Dz = torch.randn(1, 1, 1, 1)

    Dq = torch.cat((Dx, Dy, Dz), dim=1)
    Dq_norm = torch.norm(Dq)

    Dq_prox = TV.prox_l21(Dq, prox_val, dim=1)
    Dq_prox_norm = torch.norm(Dq_prox)

    torch.testing.assert_allclose(
        max(Dq_norm, 0) - prox_val,
        Dq_prox_norm,
    )
    torch.testing.assert_allclose(
        Dq / Dq_norm,
        Dq_prox / Dq_prox_norm,
    )
//  [python]
def test_prox_l21():
    prox_val = .75

    Dx = torch.randn(1, 1, 1, 1)
    Dy = torch.randn(1, 1, 1, 1)
    Dz = torch.randn(1, 1, 1, 1)

    Dq = torch.cat((Dx, Dy, Dz), dim=1)
    Dq_norm = torch.linalg.norm(Dq)

    Dq_prox = TV.prox_l21(Dq, prox_val, dim=1)
    Dq_prox_norm = torch.linalg.norm(Dq_prox)

    torch.testing.assert_allclose(
        max(Dq_norm, 0) - prox_val,
        Dq_prox_norm,
    )
    torch.testing.assert_allclose(
        Dq / Dq_norm,
        Dq_prox / Dq_prox_norm,
    )
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, geometry_hair, n_importance, inv_s, head_label, hair_label):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, p=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        
        alpha_sdf = self.compute_sdf_alpha_for_upsample(z_vals, sdf, inside_sphere, inv_s)
        valid_head_interval = (head_label[..., :-1] == head_label[..., 1:]) & (head_label[..., :-1] != -1)
        valid_hair_interval = (hair_label[..., :-1] == hair_label[..., 1:]) & (hair_label[..., :-1] != -1) 
        alpha_sdf *= valid_head_interval
        if self.blended_upsample:
            alpha_hair = self.compute_sdf_alpha_for_upsample(z_vals, geometry_hair, inside_sphere, inv_s)
            alpha_hair *= valid_hair_interval

            alpha = (alpha_sdf + alpha_hair).clamp(0, 1)
        else:
            alpha = alpha_sdf

        valid_interval = valid_head_interval | valid_hair_interval
        
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + valid_interval * 1e-7], -1), -1)[:, :-1]
        z_samples = sample_pdf(z_vals, weights, n_importance, det=True, valid_interval=valid_interval, detailed_output=True)
        return z_samples
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, geometry_hair, n_importance, inv_s, head_label, hair_label):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        
        alpha_sdf = self.compute_sdf_alpha_for_upsample(z_vals, sdf, inside_sphere, inv_s)
        valid_head_interval = (head_label[..., :-1] == head_label[..., 1:]) & (head_label[..., :-1] != -1)
        valid_hair_interval = (hair_label[..., :-1] == hair_label[..., 1:]) & (hair_label[..., :-1] != -1) 
        alpha_sdf *= valid_head_interval
        if self.blended_upsample:
            alpha_hair = self.compute_sdf_alpha_for_upsample(z_vals, geometry_hair, inside_sphere, inv_s)
            alpha_hair *= valid_hair_interval

            alpha = (alpha_sdf + alpha_hair).clamp(0, 1)
        else:
            alpha = alpha_sdf

        valid_interval = valid_head_interval | valid_hair_interval
        
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + valid_interval * 1e-7], -1), -1)[:, :-1]
        z_samples = sample_pdf(z_vals, weights, n_importance, det=True, valid_interval=valid_interval, detailed_output=True)
        return z_samples
//  [python]
def pairwise_cosine(
    x: torch.Tensor, y: torch.Tensor, *, y2: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """Compute pair-wise cosine distance between x and y.

    Parameters
    ----------
    x : torch.Tensor
        A 2-D ``[M, D]`` tensor, containing `M` vectors.
    y : torch.Tensor
        A 2-D ``[N, D]`` tensor, containing `N` vectors.

    Returns
    -------
    A ``[M, N]`` tensor with pair-wise cosine distances between x and y.
    """
    if len(x.shape) != 2 or len(y.shape) != 2:
        raise ValueError(
            f"x and y must be 2-D matrix, got: x.shape={x.shape}, y.shape={y.shape}"
        )
    if y2 is None:
        y2: torch.Tensor = torch.norm(y, dim=1)
    return _pairwise_cosine(x, y, y2)
//  [python]
def pairwise_cosine(
    x: torch.Tensor, y: torch.Tensor, *, y2: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """Compute pair-wise cosine distance between x and y.

    Parameters
    ----------
    x : torch.Tensor
        A 2-D ``[M, D]`` tensor, containing `M` vectors.
    y : torch.Tensor
        A 2-D ``[N, D]`` tensor, containing `N` vectors.

    Returns
    -------
    A ``[M, N]`` tensor with pair-wise cosine distances between x and y.
    """
    if len(x.shape) != 2 or len(y.shape) != 2:
        raise ValueError(
            f"x and y must be 2-D matrix, got: x.shape={x.shape}, y.shape={y.shape}"
        )
    if y2 is None:
        y2: torch.Tensor = torch.linalg.norm(y, dim=1)
    return _pairwise_cosine(x, y, y2)
//  [python]
def test_SACKernel(self):
        decomposition = {"1": [0, 3], "2": [1, 2]}
        kernel = SACKernel(decomposition=decomposition, batch_shape=torch.Size([]))

        self.assertIsInstance(kernel.kernel_dict, ModuleDict)
        self.assertIsInstance(kernel.base_kernel, MaternKernel)
        self.assertDictEqual(kernel.decomposition, decomposition)

        # test diag works well for lazy tensor
        x1 = torch.rand(5, 4)
        x2 = torch.rand(5, 4)
        res = kernel(x1, x2).to_dense()
        res_diag = kernel(x1, x2, diag=True)
        self.assertLess(torch.norm(res_diag - res.diag()), 1e-4)

        # test raise of ValueError
        with self.assertRaises(ValueError):
            SACKernel(decomposition={"1": [0, 3], "2": [1]}, batch_shape=torch.Size([]))
//  [python]
def test_SACKernel(self):
        decomposition = {"1": [0, 3], "2": [1, 2]}
        kernel = SACKernel(decomposition=decomposition, batch_shape=torch.Size([]))

        self.assertIsInstance(kernel.kernel_dict, ModuleDict)
        self.assertIsInstance(kernel.base_kernel, MaternKernel)
        self.assertDictEqual(kernel.decomposition, decomposition)

        # test diag works well for lazy tensor
        x1 = torch.rand(5, 4)
        x2 = torch.rand(5, 4)
        res = kernel(x1, x2).to_dense()
        res_diag = kernel(x1, x2, diag=True)
        self.assertLess(torch.linalg.norm(res_diag - res.diag()), 1e-4)

        # test raise of ValueError
        with self.assertRaises(ValueError):
            SACKernel(decomposition={"1": [0, 3], "2": [1]}, batch_shape=torch.Size([]))