jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, ref, runtime_dict=None):
        if (runtime_dict is not None) and self.enabled:
            if 'base_bxyz' in runtime_dict:
                # computing volumes
                base_bxyz = runtime_dict['base_bxyz']
                
                e_base, e_voxel, e_weight = self.voxel_graph(base_bxyz, ref)

                num_voxels = ref.bcenter.shape[0]

                ref.bxyz = scatter(base_bxyz[e_base], e_voxel, dim=0,
                                   dim_size=num_voxels, reduce='sum')
                ref.volume = scatter(torch.ones_like(e_base), e_voxel, dim=0,
                                     dim_size=num_voxels, reduce='sum')
                mask = (ref.volume > 0.5)
                ref.volume_mask = mask

                ref.bxyz[mask] = ref.bxyz[mask] / ref.volume[mask, None]
                ref.bxyz[~mask] = ref.bcenter[~mask]

                point_d = base_bxyz[e_base, 1:] - ref.bxyz[e_voxel, 1:]
                point_ddT = point_d.unsqueeze(-1) * point_d.unsqueeze(-2)
                voxel_ddT = scatter(point_ddT, e_voxel, dim=0,
                                    dim_size=num_voxels, reduce='mean')

                if True:
                    voxel_eigvals, voxel_eigvecs = torch.linalg.eigh(voxel_ddT)
                    #voxel_eigvals, voxel_eigvecs = np.linalg.eigh(voxel_ddT.detach().cpu().numpy())
                    #voxel_eigvals = torch.from_numpy(voxel_eigvals).to(voxel_ddT)
                    #voxel_eigvecs = torch.from_numpy(voxel_eigvecs).to(voxel_ddT)
                    ref.eigvals = voxel_eigvals
                    ref.eigvecs = voxel_eigvecs
                    point_l1_proj = (point_d.unsqueeze(-2) @ ref.eigvecs[e_voxel])[:, 0]
                    ref.l1_proj_max = scatter(point_l1_proj, e_voxel, dim=0,
                                              dim_size=num_voxels, reduce='max')
                    ref.l1_proj_min = scatter(point_l1_proj, e_voxel, dim=0,
                                              dim_size=num_voxels, reduce='min')

                #import polyscope as ps; ps.init(); ps.set_up_dir('z_up')

                #ps.register_point_cloud('voxel_centers', ref.bcenter[:, 1:].detach().cpu().numpy(), radius=2e-4)
                #ps.register_point_cloud('base', base_bxyz[:, 1:].detach().cpu().numpy(), radius=4e-5)
                #centers = torch.cat([ref.bcenter, base_bxyz], dim=0)[:, 1:].detach().cpu().numpy()
                #edges = torch.stack([e_base+num_voxels, e_voxel], dim=-1).detach().cpu().numpy()
                #ps.register_curve_network('edges', centers, edges, radius=2e-5)
                #ps.show()
                #import ipdb; ipdb.set_trace()
        
        return ref
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        if ODhat.numel() == 0:
            raise FullyTransparentException("ODhat is empty")

        # compute eigenvectors
        try:
            _, eigvecs = torch.symeig(cov(ODhat.T))
        except Exception as e:
            if "LinAlgError" in e.__class__.__name__:  # should be torch._C._LinalgError but that's not exposed
                raise FullyTransparentException("could not compute eigenvectors")
            raise e
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        if ODhat.numel() == 0:
            raise FullyTransparentException("ODhat is empty")

        # compute eigenvectors
        try:
            _, eigvecs = torch.linalg.eigh(cov(ODhat.T))
        except Exception as e:
            if "LinAlgError" in e.__class__.__name__:  # should be torch._C._LinalgError but that's not exposed
                raise FullyTransparentException("could not compute eigenvectors")
            raise e
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def _point_normals(
    self,
    *,
    scale: Number | None = None,
    **kwargs,
) -> Points:
    if scale is None:
        if self.triangles is None:
            msg = "If no triangles are provided, you must specify a scale."
            raise ValueError(msg)

        tri_n = self.triangle_normals  # N.B.: magnitude = triangle area
        n = torch.zeros_like(self.points)
        # TODO: instead of distributing to the vertices equally, we should
        #       distribute according to the angles of the triangles.
        for k in range(3):
            # n[self.triangles[k, i]] += tri_n[i]
            n.scatter_reduce_(
                dim=0,
                index=self.triangles[k].view(-1, 1),
                src=tri_n,
                reduce="sum",
            )

    else:
        # Get a smooth field of normals via the Structure Tensor:
        local_cov = self.point_moments(
            order=2, scale=scale, central=True, **kwargs
        )
        local_QL = torch.symeig(local_cov)
        local_nuv = local_QL.eigenvectors  # (N, 3, 3)
        n = local_nuv[:, :, 0]

        # Orient the normals according to the triangles, if any:
        if self.triangles is not None:
            n_0 = self.point_normals(scale=None, **kwargs)
            assert n_0.shape == (self.n_points, 3)

        else:
            n_0 = n.mean(dim=0, keepdim=True)
            assert n_0.shape == (1, 3)

        n = (n_0 * n).sum(-1).sign().view(-1, 1) * n

    # Try to enforce some consistency...
    if False:
        n = F.normalize(n, p=2, dim=-1)
        # n_e = n[self.edges[0]] + n[self.edges[1]]
        # The backward of torch.index_select is much faster than that of
        # indexing:
        n_e = torch.index_select(n, 0, self.edges[0]) + torch.index_select(
            n, 0, self.edges[1]
        )
        assert n_e.shape == (len(self.edges[0]), 3)

        n_v = torch.zeros_like(n)
        n_v.scatter_reduce_(
            dim=0,
            index=self.edges.reshape(-1, 1),
            src=n_e.tile(2, 1),
            reduce="mean",
            include_self=False,
        )

        assert n_v.shape == (self.n_points, 3)
        # n_v = n_v - n
        n = (n_v * n).sum(-1).sign().view(-1, 1) * n

    n = F.normalize(n, p=2, dim=-1)
    assert n.shape == (self.n_points, 3)
    return n
//  [python]
def _point_normals(
    self,
    *,
    scale: Number | None = None,
    **kwargs,
) -> Points:
    if scale is None:
        if self.triangles is None:
            msg = "If no triangles are provided, you must specify a scale."
            raise ValueError(msg)

        tri_n = self.triangle_normals  # N.B.: magnitude = triangle area
        n = torch.zeros_like(self.points)
        # TODO: instead of distributing to the vertices equally, we should
        #       distribute according to the angles of the triangles.
        for k in range(3):
            # n[self.triangles[k, i]] += tri_n[i]
            n.scatter_reduce_(
                dim=0,
                index=self.triangles[k].view(-1, 1),
                src=tri_n,
                reduce="sum",
            )

    else:
        # Get a smooth field of normals via the Structure Tensor:
        local_cov = self.point_moments(
            order=2, scale=scale, central=True, **kwargs
        )
        local_QL = torch.linalg.eigh(local_cov)
        local_nuv = local_QL.eigenvectors  # (N, 3, 3)
        n = local_nuv[:, :, 0]

        # Orient the normals according to the triangles, if any:
        if self.triangles is not None:
            n_0 = self.point_normals(scale=None, **kwargs)
            assert n_0.shape == (self.n_points, 3)

        else:
            n_0 = n.mean(dim=0, keepdim=True)
            assert n_0.shape == (1, 3)

        n = (n_0 * n).sum(-1).sign().view(-1, 1) * n

    # Try to enforce some consistency...
    if False:
        n = F.normalize(n, p=2, dim=-1)
        # n_e = n[self.edges[0]] + n[self.edges[1]]
        # The backward of torch.index_select is much faster than that of
        # indexing:
        n_e = torch.index_select(n, 0, self.edges[0]) + torch.index_select(
            n, 0, self.edges[1]
        )
        assert n_e.shape == (len(self.edges[0]), 3)

        n_v = torch.zeros_like(n)
        n_v.scatter_reduce_(
            dim=0,
            index=self.edges.reshape(-1, 1),
            src=n_e.tile(2, 1),
            reduce="mean",
            include_self=False,
        )

        assert n_v.shape == (self.n_points, 3)
        # n_v = n_v - n
        n = (n_v * n).sum(-1).sign().view(-1, 1) * n

    n = F.normalize(n, p=2, dim=-1)
    assert n.shape == (self.n_points, 3)
    return n
//  [python]
def plot_multivariate_gaussian(mu, Sigma, sigma_range=3.0):
  device = mu.device

  e,Q = torch.symeig(Sigma)

  Sigma_sqrt = Q @ torch.diag(torch.sqrt(e)) @ Q.T

  N = 1001
  sigma_range = torch.linspace(-sigma_range, sigma_range, N, device=device)
  x, y = torch.meshgrid(sigma_range, sigma_range, indexing='ij')

  pos = torch.empty(x.shape + (2,), device=device)
  pos[:, :, 0] = x
  pos[:, :, 1] = y
  
  transformed_pos = torch.squeeze(Sigma_sqrt[None,None,:] @ pos[:,:,:,None])
  transformed_pos[:,:,None] += mu

  n = mu.shape[0]
  Sigma_det = torch.det(Sigma)
  Sigma_inv = torch.inverse(Sigma) 

  norm_constant = torch.sqrt((2*np.pi)**n * Sigma_det)
  fac = torch.einsum('...k,kl,...l->...', transformed_pos-mu, Sigma_inv, transformed_pos-mu)

  x_t = transformed_pos[:,:,0]
  y_t = transformed_pos[:,:,1]
  z = torch.exp(-fac / 2) / norm_constant

  max_border = max(torch.max(z[N-1,:]), torch.max(z[:,N-1]))
  z = np.ma.array(z, mask=z < max_border)

  return x_t, y_t, z
//  [python]
def plot_multivariate_gaussian(mu, Sigma, sigma_range=3.0):
  device = mu.device

  e,Q = torch.linalg.eigh(Sigma)

  Sigma_sqrt = Q @ torch.diag(torch.sqrt(e)) @ Q.T

  N = 1001
  sigma_range = torch.linspace(-sigma_range, sigma_range, N, device=device)
  x, y = torch.meshgrid(sigma_range, sigma_range, indexing='ij')

  pos = torch.empty(x.shape + (2,), device=device)
  pos[:, :, 0] = x
  pos[:, :, 1] = y
  
  transformed_pos = torch.squeeze(Sigma_sqrt[None,None,:] @ pos[:,:,:,None])
  transformed_pos[:,:,None] += mu

  n = mu.shape[0]
  Sigma_det = torch.linalg.det(Sigma)
  Sigma_inv = torch.linalg.inv(Sigma) 

  norm_constant = torch.sqrt((2*np.pi)**n * Sigma_det)
  fac = torch.einsum('...k,kl,...l->...', transformed_pos-mu, Sigma_inv, transformed_pos-mu)

  x_t = transformed_pos[:,:,0]
  y_t = transformed_pos[:,:,1]
  z = torch.exp(-fac / 2) / norm_constant

  max_border = max(torch.max(z[N-1,:]), torch.max(z[:,N-1]))
  z = np.ma.array(z, mask=z < max_border)

  return x_t, y_t, z
//  [python]
def matrix_sqrth(A_NFMM):
    eig_val_NFM, eig_vec_NFMM = torch.symeig(A_NFMM)
    eig_val_NFM[eig_val_NFM < EPS] = EPS
    return torch.einsum("nfij, nfj, nflj -> nfil", eig_vec_NFMM, torch.sqrt(eig_val_NFM), eig_vec_NFMM.conj())
//  [python]
def matrix_sqrth(A_NFMM):
    eig_val_NFM, eig_vec_NFMM = torch.linalg.eigh(A_NFMM)
    eig_val_NFM[eig_val_NFM < EPS] = EPS
    return torch.einsum("nfij, nfj, nflj -> nfil", eig_vec_NFMM, torch.sqrt(eig_val_NFM), eig_vec_NFMM.conj())
//  [python]
def _step(
    gradient: torch.Tensor,
    hessian: torch.Tensor,
    trust_radius: torch.Tensor,
    config: LevenbergMarquardtConfig,
) -> tuple[torch.Tensor, torch.Tensor, bool, torch.Tensor]:
    """Compute the next Levenberg–Marquardt step.

    Args:
        gradient: The gradient with ``shape=(n,)``.
        hessian: The hessian with ``shape=(n, n)``.
        trust_radius: The target trust radius.
        config: The optimizer config.

    Notes:
        * the code to 'excise' certain parameters is for now removed until its clear
          it is needed.
        * only trust region is implemented (i.e., only trust0 > 0 is supported)

    Returns:
        The step with ``shape=(n,)``, the expected improvement with ``shape=()``,
        a boolean indicating whether the damping factor was adjusted, and the damping
        factor with ``shape=(1,)``.
    """
    from scipy import optimize

    eigenvalues, _ = torch.symeig(hessian)
    eigenvalue_smallest = eigenvalues.min()

    if eigenvalue_smallest < config.min_eigenvalue:
        # Mix in SD step if Hessian minimum eigenvalue is negative - experimental.
        adjacency = (
            max(config.min_eigenvalue, 0.01 * abs(eigenvalue_smallest))
            - eigenvalue_smallest
        )

        _LOGGER.info(
            f"hessian has a small or negative eigenvalue ({eigenvalue_smallest:.1e}), "
            f"mixing in some steepest descent ({adjacency:.1e}) to correct this."
        )
        hessian += adjacency * torch.eye(
            hessian.shape[0], device=hessian.device, dtype=hessian.dtype
        )

    damping_factor = torch.tensor(1.0)

    dx, improvement = _solver(damping_factor, gradient, hessian)
    dx_norm = torch.linalg.norm(dx)

    adjust_damping = bool(dx_norm > trust_radius)

    if adjust_damping:
        # LPW tried a few optimizers and found Brent works well, but also that the
        # tolerance is fractional - if the optimized value is zero it takes a lot of
        # meaningless steps.
        damping_factor = optimize.brent(
            _damping_factor_loss_fn,
            (
                gradient.detach().cpu(),
                hessian.detach().cpu(),
                trust_radius.detach().cpu(),
            ),
            brack=(config.min_damping_factor, config.min_damping_factor * 4),
            tol=1.0e-6 if config.mode.lower() == _ADAPTIVE else 1.0e-4,
        )

        dx, improvement = _solver(damping_factor, gradient, hessian)

    return dx, improvement, adjust_damping, damping_factor
//  [python]
def _step(
    gradient: torch.Tensor,
    hessian: torch.Tensor,
    trust_radius: torch.Tensor,
    config: LevenbergMarquardtConfig,
) -> tuple[torch.Tensor, torch.Tensor, bool, torch.Tensor]:
    """Compute the next Levenberg–Marquardt step.

    Args:
        gradient: The gradient with ``shape=(n,)``.
        hessian: The hessian with ``shape=(n, n)``.
        trust_radius: The target trust radius.
        config: The optimizer config.

    Notes:
        * the code to 'excise' certain parameters is for now removed until its clear
          it is needed.
        * only trust region is implemented (i.e., only trust0 > 0 is supported)

    Returns:
        The step with ``shape=(n,)``, the expected improvement with ``shape=()``,
        a boolean indicating whether the damping factor was adjusted, and the damping
        factor with ``shape=(1,)``.
    """
    from scipy import optimize

    eigenvalues, _ = torch.linalg.eigh(hessian)
    eigenvalue_smallest = eigenvalues.min()

    if eigenvalue_smallest < config.min_eigenvalue:
        # Mix in SD step if Hessian minimum eigenvalue is negative - experimental.
        adjacency = (
            max(config.min_eigenvalue, 0.01 * abs(eigenvalue_smallest))
            - eigenvalue_smallest
        )

        _LOGGER.info(
            f"hessian has a small or negative eigenvalue ({eigenvalue_smallest:.1e}), "
            f"mixing in some steepest descent ({adjacency:.1e}) to correct this."
        )
        hessian += adjacency * torch.eye(
            hessian.shape[0], device=hessian.device, dtype=hessian.dtype
        )

    damping_factor = torch.tensor(1.0)

    dx, improvement = _solver(damping_factor, gradient, hessian)
    dx_norm = torch.linalg.norm(dx)

    adjust_damping = bool(dx_norm > trust_radius)

    if adjust_damping:
        # LPW tried a few optimizers and found Brent works well, but also that the
        # tolerance is fractional - if the optimized value is zero it takes a lot of
        # meaningless steps.
        damping_factor = optimize.brent(
            _damping_factor_loss_fn,
            (
                gradient.detach().cpu(),
                hessian.detach().cpu(),
                trust_radius.detach().cpu(),
            ),
            brack=(config.min_damping_factor, config.min_damping_factor * 4),
            tol=1.0e-6 if config.mode.lower() == _ADAPTIVE else 1.0e-4,
        )

        dx, improvement = _solver(damping_factor, gradient, hessian)

    return dx, improvement, adjust_damping, damping_factor
//  [python]
def score_K(K, q=1, p=None, normalize=False):
    if normalize:
        K = normalize_K(K)
    K_ = weight_K(K, p)
    w, _ = torch.symeig(K_)
    return torch.exp(entropy_q(w, q=q))
//  [python]
def score_K(K, q=1, p=None, normalize=False):
    if normalize:
        K = normalize_K(K)
    K_ = weight_K(K, p)
    w, _ = torch.linalg.eigh(K_)
    return torch.exp(entropy_q(w, q=q))
//  [python]
def get_electrostatic_function_encoding(edge_index, num_nodes):
    """Kernel based on the electrostatic interaction between nodes.
    """
    L = to_scipy_sparse_matrix(
        *get_laplacian(edge_index, normalization=None, num_nodes=num_nodes)
    ).todense()
    L = torch.as_tensor(L)
    Dinv = torch.eye(L.shape[0]) * (L.diag() ** -1)
    A = deepcopy(L).abs()
    A.fill_diagonal_(0)
    DinvA = Dinv.matmul(A)

    evals, evecs = torch.symeig(L, eigenvectors=True)
    offset = (evals < EPS).sum().item()
    if offset == num_nodes:
        return torch.zeros(num_nodes, 8, dtype=torch.float32)

    electrostatic = evecs[:, offset:] / evals[offset:] @ evecs[:, offset:].T
    electrostatic = electrostatic - electrostatic.diag()
    green_encoding = torch.stack([
        electrostatic.min(dim=0)[0],  # Min of Vi -> j
        electrostatic.mean(dim=0),  # Mean of Vi -> j
        electrostatic.std(dim=0),  # Std of Vi -> j
        electrostatic.min(dim=1)[0],  # Min of Vj -> i
        electrostatic.std(dim=1),  # Std of Vj -> i
        (DinvA * electrostatic).sum(dim=0),  # Mean of interaction on direct neighbour
        (DinvA * electrostatic).sum(dim=1),  # Mean of interaction from direct neighbour
    ], dim=1)

    return green_encoding
//  [python]
def get_electrostatic_function_encoding(edge_index, num_nodes):
    """Kernel based on the electrostatic interaction between nodes.
    """
    L = to_scipy_sparse_matrix(
        *get_laplacian(edge_index, normalization=None, num_nodes=num_nodes)
    ).todense()
    L = torch.as_tensor(L)
    Dinv = torch.eye(L.shape[0]) * (L.diag() ** -1)
    A = deepcopy(L).abs()
    A.fill_diagonal_(0)
    DinvA = Dinv.matmul(A)

    evals, evecs = torch.linalg.eigh(L)
    offset = (evals < EPS).sum().item()
    if offset == num_nodes:
        return torch.zeros(num_nodes, 8, dtype=torch.float32)

    electrostatic = evecs[:, offset:] / evals[offset:] @ evecs[:, offset:].T
    electrostatic = electrostatic - electrostatic.diag()
    green_encoding = torch.stack([
        electrostatic.min(dim=0)[0],  # Min of Vi -> j
        electrostatic.mean(dim=0),  # Mean of Vi -> j
        electrostatic.std(dim=0),  # Std of Vi -> j
        electrostatic.min(dim=1)[0],  # Min of Vj -> i
        electrostatic.std(dim=1),  # Std of Vj -> i
        (DinvA * electrostatic).sum(dim=0),  # Mean of interaction on direct neighbour
        (DinvA * electrostatic).sum(dim=1),  # Mean of interaction from direct neighbour
    ], dim=1)

    return green_encoding
//  [python]
def get_rotation_to_best_fit_xy(
    points: torch.Tensor, centroid: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """
    Returns a rotation R such that `points @ R` has a best fit plane
    parallel to the xy plane

    Args:
        points: (*, N, 3) tensor of points in 3D
        centroid: (*, 1, 3), (3,) or scalar: their centroid

    Returns:
        (*, 3, 3) tensor rotation matrix
    """
    if centroid is None:
        centroid = points.mean(dim=-2, keepdim=True)

    points_centered = points - centroid
    e, v = torch.symeig(points_centered.transpose(-1, -2) @ points_centered, eigenvectors=True)
    # in general, evec can form either right- or left-handed basis,
    # but we need the former to have a proper rotation (not reflection)
    return torch.cat(
        (v[..., 1:], torch.cross(v[..., 1], v[..., 2])[..., None]), dim=-1
    )
//  [python]
def get_rotation_to_best_fit_xy(
    points: torch.Tensor, centroid: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """
    Returns a rotation R such that `points @ R` has a best fit plane
    parallel to the xy plane

    Args:
        points: (*, N, 3) tensor of points in 3D
        centroid: (*, 1, 3), (3,) or scalar: their centroid

    Returns:
        (*, 3, 3) tensor rotation matrix
    """
    if centroid is None:
        centroid = points.mean(dim=-2, keepdim=True)

    points_centered = points - centroid
    _, evec = torch.linalg.eigh(points_centered.transpose(-1, -2) @ points_centered)
    # in general, evec can form either right- or left-handed basis,
    # but we need the former to have a proper rotation (not reflection)
    return torch.cat(
        (evec[..., 1:], torch.cross(evec[..., 1], evec[..., 2])[..., None]), dim=-1
    )
//  [python]
def test_lkj_prior_sample(self):
        prior = LKJCovariancePrior(2, 0.5, sd_prior=SmoothedBoxPrior(exp(-1), exp(1)))
        random_samples = prior.sample(torch.Size((6,)))
        # need to check that these are positive semi-sefinite
        min_eval = torch.symeig(random_samples, eigenvectors=False)[0].min()
        self.assertTrue(min_eval >= 0)
        # and that they are symmetric
        max_non_symm = (random_samples - random_samples.transpose(-1, -2)).abs().max()
        self.assertLess(max_non_symm, 1e-4)

        self.assertEqual(random_samples.shape, torch.Size((6, 2, 2)))
//  [python]
def test_lkj_prior_sample(self):
        prior = LKJCovariancePrior(2, 0.5, sd_prior=SmoothedBoxPrior(exp(-1), exp(1)))
        random_samples = prior.sample(torch.Size((6,)))
        # need to check that these are positive semi-sefinite
        min_eval = torch.linalg.eigh(random_samples)[0].min()
        self.assertTrue(min_eval >= 0)
        # and that they are symmetric
        max_non_symm = (random_samples - random_samples.transpose(-1, -2)).abs().max()
        self.assertLess(max_non_symm, 1e-4)

        self.assertEqual(random_samples.shape, torch.Size((6, 2, 2)))
//  [python]
def EVD_Laplacian(data, norm=None):
    L_raw = get_laplacian(to_undirected(data.edge_index, num_nodes=data.num_nodes), 
                          normalization=norm, num_nodes=data.num_nodes)
    L = SparseTensor(row=L_raw[0][0], col=L_raw[0][1], value=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    D, V  = torch.symeig(L)
    return D, V
//  [python]
def EVD_Laplacian(data, norm=None):
    L_raw = get_laplacian(to_undirected(data.edge_index, num_nodes=data.num_nodes), 
                          normalization=norm, num_nodes=data.num_nodes)
    L = SparseTensor(row=L_raw[0][0], col=L_raw[0][1], value=L_raw[1], sparse_sizes=(data.num_nodes, data.num_nodes)).to_dense()

    D, V  = torch.linalg.eigh(L)
    return D, V
//  [python]
def _fit(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the principal components of the reference data.

        We compute the principal components of the reference data using the covariance matrix and then
        remove the largest `n_components` eigenvectors. The remaining eigenvectors correspond to the
        invariant dimensions of the data. Changes in these dimensions are used to compute the outlier
        score which is the distance to the principal subspace spanned by the first `n_components`
        eigenvectors.

        Parameters
        ----------
        x
            The reference data.

        Returns
        -------
        The principal components of the reference data.

        Raises
        ------
        ValueError
            If `n_components` is greater than or equal to number of features
        """
        if self.n_components >= x.shape[1]:
            raise ValueError("n_components must be less than the number of features.")

        self.x_ref_mean = x.mean(0)
        x -= self.x_ref_mean
        cov_mat = (x.t() @ x)/(len(x)-1)
        e, V = torch.symeig(cov_mat, eigenvectors=True)
        return V[:, :-self.n_components]
//  [python]
def _fit(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the principal components of the reference data.

        We compute the principal components of the reference data using the covariance matrix and then
        remove the largest `n_components` eigenvectors. The remaining eigenvectors correspond to the
        invariant dimensions of the data. Changes in these dimensions are used to compute the outlier
        score which is the distance to the principal subspace spanned by the first `n_components`
        eigenvectors.

        Parameters
        ----------
        x
            The reference data.

        Returns
        -------
        The principal components of the reference data.

        Raises
        ------
        ValueError
            If `n_components` is greater than or equal to number of features
        """
        if self.n_components >= x.shape[1]:
            raise ValueError("n_components must be less than the number of features.")

        self.x_ref_mean = x.mean(0)
        x -= self.x_ref_mean
        cov_mat = (x.t() @ x)/(len(x)-1)
        _, V = torch.linalg.eigh(cov_mat)
        return V[:, :-self.n_components]
//  [python]
def _fit(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the principal components of the reference data.

        We compute the principal components of the reference data using the kernel matrix and then
        return the largest `n_components` eigenvectors. These are then normalized to have length
        equal to `1/eigenvalue`. Note that this differs from the linear case where we remove the
        largest eigenvectors.

        Parameters
        ----------
        x
            The reference data.

        Returns
        -------
        The principal components of the reference data.

        Raises
        ------
        ValueError
            If `n_components` is greater than or equal to the number of reference samples.
        """
        if self.n_components >= x.shape[0]:
            raise ValueError("n_components must be less than the number of reference instances.")

        self.x_ref = x
        K = self.compute_kernel_mat(x)
        D, V = torch.symeig(K, eigenvectors=True)
        pcs = V / torch.sqrt(D)[None, :]
        return pcs[:, -self.n_components:]
//  [python]
def _fit(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the principal components of the reference data.

        We compute the principal components of the reference data using the kernel matrix and then
        return the largest `n_components` eigenvectors. These are then normalized to have length
        equal to `1/eigenvalue`. Note that this differs from the linear case where we remove the
        largest eigenvectors.

        Parameters
        ----------
        x
            The reference data.

        Returns
        -------
        The principal components of the reference data.

        Raises
        ------
        ValueError
            If `n_components` is greater than or equal to the number of reference samples.
        """
        if self.n_components >= x.shape[0]:
            raise ValueError("n_components must be less than the number of reference instances.")

        self.x_ref = x
        K = self.compute_kernel_mat(x)
        D, V = torch.linalg.eigh(K)
        pcs = V / torch.sqrt(D)[None, :]
        return pcs[:, -self.n_components:]
//  [python]
def eigenpairs(L: B.TorchNumeric, k: int):
    """
    Obtain the eigenpairs that correspond to the `k` lowest eigenvalues
    of a symmetric positive semi-definite matrix `L`.

    TODO(AR): Replace with torch.lobpcg after sparse matrices are supported
    by torch.
    """
    l, u = torch.symeig(L, eigenvectors=True)
    return l[:k], u[:, :k]
//  [python]
def eigenpairs(L: B.TorchNumeric, k: int):
    """
    Obtain the eigenpairs that correspond to the `k` lowest eigenvalues
    of a symmetric positive semi-definite matrix `L`.

    TODO(AR): Replace with torch.lobpcg after sparse matrices are supported
    by torch.
    """
    l, u = torch.linalg.eigh(L)
    return l[:k], u[:, :k]
//  [python]
def _extract_shape(idx, x, j, cur_center):
    _a = []
    for i in range(len(idx)):
        if idx[i] == j:
            if torch.sum(cur_center)==0:
                opt_x = x[i]
            else:
                opt_x = _sbd(cur_center, x[i])
            _a.append(opt_x)
            
    if len(_a) == 0:
        indices = torch.randperm(x.shape[0])[:1]
        return torch.squeeze(x[indices].clone())
        #return torch.zeros((x.shape[1]))

    a = torch.stack(_a)
    
    columns = a.shape[1]
    y = zscore(a, axis=1, ddof=1)
    
    s = y[:, :, 0].transpose(0, 1).mm(y[:, :, 0])

    p = torch.empty((columns, columns), device="cuda", dtype=torch.float32)
    p.fill_(1.0 / columns)
    p = torch.eye(columns, device="cuda", dtype=torch.float32) - p

    m = p.mm(s).mm(p)
    _, vec = torch.symeig(m, eigenvectors=True)
    centroid = vec[:, -1]

    finddistance1 = torch.norm(a.sub(centroid.reshape((x.shape[1], 1))), 2, dim=(1, 2)).sum()
    finddistance2 = torch.norm(a.add(centroid.reshape((x.shape[1], 1))), 2, dim=(1, 2)).sum()

    if finddistance1 >= finddistance2:
        centroid.mul_(-1)

    return zscore(centroid, ddof=1)
//  [python]
def _extract_shape(idx, x, j, cur_center):
    _a = []
    for i in range(len(idx)):
        if idx[i] == j:
            if torch.sum(cur_center)==0:
                opt_x = x[i]
            else:
                opt_x = _sbd(cur_center, x[i])
            _a.append(opt_x)
            
    if len(_a) == 0:
        indices = torch.randperm(x.shape[0])[:1]
        return torch.squeeze(x[indices].clone())
        #return torch.zeros((x.shape[1]))

    a = torch.stack(_a)
    
    columns = a.shape[1]
    y = zscore(a, axis=1, ddof=1)
    
    s = y[:, :, 0].transpose(0, 1).mm(y[:, :, 0])

    p = torch.empty((columns, columns), device="cuda", dtype=torch.float32)
    p.fill_(1.0 / columns)
    p = torch.eye(columns, device="cuda", dtype=torch.float32) - p

    m = p.mm(s).mm(p)
    _, vec = torch.linalg.eigh(m,UPLO='U')
    centroid = vec[:, -1]

    finddistance1 = torch.norm(a.sub(centroid.reshape((x.shape[1], 1))), 2, dim=(1, 2)).sum()
    finddistance2 = torch.norm(a.add(centroid.reshape((x.shape[1], 1))), 2, dim=(1, 2)).sum()

    if finddistance1 >= finddistance2:
        centroid.mul_(-1)

    return zscore(centroid, ddof=1)
//  [python]
def forward(self, x: torch.Tensor):
        B, N, D = x.shape
        assert D == self.n_point_dim

        device = x.device

        x_centered = x - x.mean(dim=1, keepdims=True)
        cov = torch.bmm(x_centered.transpose(2, 1) / np.sqrt(N), x_centered / np.sqrt(N))

        cov = cov.cpu()

        _, evecs = torch.symeig(cov, eigenvectors=True)
        evecs = evecs.to(device)

        if self.random_sign:
            signs = torch.randint(0, 2, (B, 1, D)).float().to(device)
            signs[signs == 0] = -1
            evecs *= signs

        if self.random_order:
            for b in range(B):
                evecs[b] = evecs[b, :, torch.randperm(D)]

        x = torch.bmm(x_centered, evecs)

        return self.base_model(x)
//  [python]
def forward(self, x: torch.Tensor):
        B, N, D = x.shape
        assert D == self.n_point_dim

        device = x.device

        x_centered = x - x.mean(dim=1, keepdims=True)
        cov = torch.bmm(x_centered.transpose(2, 1) / np.sqrt(N), x_centered / np.sqrt(N))

        cov = cov.cpu()

        _, evecs = torch.linalg.eigh(cov)
        evecs = evecs.to(device)

        if self.random_sign:
            signs = torch.randint(0, 2, (B, 1, D)).float().to(device)
            signs[signs == 0] = -1
            evecs *= signs

        if self.random_order:
            for b in range(B):
                evecs[b] = evecs[b, :, torch.randperm(D)]

        x = torch.bmm(x_centered, evecs)

        return self.base_model(x)
//  [python]
def get_eigs(w):
    """
    Computes eigenvalues of the weights.
    :param w: weights
    :return: eigenvalues
    """
    w = weights2mat(w)[0]
    evals, V = torch.symeig(w.t() @ w)
    return evals
//  [python]
def get_eigs(w):
    """
    Computes eigenvalues of the weights.
    :param w: weights
    :return: eigenvalues
    """
    w = weights2mat(w)[0]
    evals, V = torch.linalg.eigh(w.t() @ w)
    return evals
//  [python]
def compute_proj_mat(self,lamb=None):
        if self.verbose:
            print("Computing Projector Matrix for Reduced Rank Regression")

        self.lamb = self.lamb if lamb is None else lamb

        self.proj_mat = self.Vt_plus.T @ self.Vt_minus @ torch.linalg.inv(self.Vt_minus.T @ self.Vt_minus + 
                                                                          self.lamb*torch.eye(self.Vt_minus.shape[1]).to(self.device)) @ \
                                                                          self.Vt_minus.T @ self.Vt_plus
        
        self.proj_mat_S, self.proj_mat_V = torch.symeig(self.proj_mat)
        #todo: more efficient to flip ranks (negative index) in compute_reduced_rank_regression but also less interpretable
        self.proj_mat_S = torch.flip(self.proj_mat_S, dims=(0,))
        self.proj_mat_V = torch.flip(self.proj_mat_V, dims=(1,))

        if self.verbose:
            print("Projector Matrix computed! \n")
//  [python]
def compute_proj_mat(self,lamb=None):
        if self.verbose:
            print("Computing Projector Matrix for Reduced Rank Regression")

        self.lamb = self.lamb if lamb is None else lamb

        self.proj_mat = self.Vt_plus.T @ self.Vt_minus @ torch.linalg.inv(self.Vt_minus.T @ self.Vt_minus + 
                                                                          self.lamb*torch.eye(self.Vt_minus.shape[1]).to(self.device)) @ \
                                                                          self.Vt_minus.T @ self.Vt_plus
        
        self.proj_mat_S, self.proj_mat_V = torch.linalg.eigh(self.proj_mat)
        #todo: more efficient to flip ranks (negative index) in compute_reduced_rank_regression but also less interpretable
        self.proj_mat_S = torch.flip(self.proj_mat_S, dims=(0,))
        self.proj_mat_V = torch.flip(self.proj_mat_V, dims=(1,))

        if self.verbose:
            print("Projector Matrix computed! \n")
//  [python]
t0 = time()
for i in range(75):
    with torch.no_grad():
        scores = scorer.score(G.visualize(torch.tensor(codes, dtype=torch.float32, device="cuda"))) #.reshape([-1,4096])
    newcodes = optim.step_simple(scores, codes)
    codes = newcodes
Atsr = torch.tensor(optim.A).cuda()
Ctsr = Atsr @ Atsr.T
#%%
eigvals, eigvecs = torch.symeig(Ctsr)
//  [python]
t0 = time()
for i in range(75):
    with torch.no_grad():
        scores = scorer.score(G.visualize(torch.tensor(codes, dtype=torch.float32, device="cuda"))) #.reshape([-1,4096])
    newcodes = optim.step_simple(scores, codes)
    codes = newcodes
Atsr = torch.tensor(optim.A).cuda()
Ctsr = Atsr @ Atsr.T
#%%
eigvals, eigvecs = torch.linalg.eigh(Ctsr)
//  [python]
def forward(self, x: torch.Tensor):
        B, N, D = x.shape
        assert D == self.n_point_dim

        device = x.device

        x_centered = x - x.mean(dim=1, keepdims=True)
        cov = torch.bmm(x_centered.transpose(2, 1) / np.sqrt(N), x_centered / np.sqrt(N))

        cov = cov.cpu()

        evals, evecs = torch.symeig(cov, eigenvectors=True)  # B x D x D

        order_ambiguity_mask = torch.zeros(B, dtype=torch.bool)

        for i, j in itertools.combinations(range(D), 2):
            order_ambiguity_mask |= torch.isclose(evals[:, i], evals[:, j])

        if not self.stack_ambiguous_orders:
            order_ambiguity_mask[:] = False

        if order_ambiguity_mask.sum() < B:
            stacked_poses = self._stack_sign_ambiguities(
                                evecs[~order_ambiguity_mask]).to(device)

            x = torch.bmm(x_centered[~order_ambiguity_mask], stacked_poses)  # B x N x (P * d)
            x = x.view(B, N, -1, D)  # B x N x P x D

            logits_no_order_ambiguity, _ = self.base_model(x)

            n_classes = logits_no_order_ambiguity.shape[-1]

        if order_ambiguity_mask.sum() > 0:
            stacked_poses = self._stack_order_and_sign_ambiguities(
                                    evecs[order_ambiguity_mask]).to(device)

            x = torch.bmm(x_centered[order_ambiguity_mask], stacked_poses)  # B x N x (P * d)
            x = x.view(B, N, -1, D)  # B x N x P x D

            logits_order_ambiguity, _ = self.base_model(x)

            n_classes = logits_order_ambiguity.shape[-1]

        logits = torch.zeros((B, n_classes)).to(device)

        if order_ambiguity_mask.sum() < B:
            logits[~order_ambiguity_mask] = logits_no_order_ambiguity

        if order_ambiguity_mask.sum() > 0:
            logits[order_ambiguity_mask] = logits_order_ambiguity

        return logits, None
//  [python]
def forward(self, x: torch.Tensor):
        B, N, D = x.shape
        assert D == self.n_point_dim

        device = x.device

        x_centered = x - x.mean(dim=1, keepdims=True)
        cov = torch.bmm(x_centered.transpose(2, 1) / np.sqrt(N), x_centered / np.sqrt(N))

        cov = cov.cpu()

        evals, evecs = torch.linalg.eigh(cov)  # B x D x D

        order_ambiguity_mask = torch.zeros(B, dtype=torch.bool)

        for i, j in itertools.combinations(range(D), 2):
            order_ambiguity_mask |= torch.isclose(evals[:, i], evals[:, j])

        if not self.stack_ambiguous_orders:
            order_ambiguity_mask[:] = False

        if order_ambiguity_mask.sum() < B:
            stacked_poses = self._stack_sign_ambiguities(
                                evecs[~order_ambiguity_mask]).to(device)

            x = torch.bmm(x_centered[~order_ambiguity_mask], stacked_poses)  # B x N x (P * d)
            x = x.view(B, N, -1, D)  # B x N x P x D

            logits_no_order_ambiguity, _ = self.base_model(x)

            n_classes = logits_no_order_ambiguity.shape[-1]

        if order_ambiguity_mask.sum() > 0:
            stacked_poses = self._stack_order_and_sign_ambiguities(
                                    evecs[order_ambiguity_mask]).to(device)

            x = torch.bmm(x_centered[order_ambiguity_mask], stacked_poses)  # B x N x (P * d)
            x = x.view(B, N, -1, D)  # B x N x P x D

            logits_order_ambiguity, _ = self.base_model(x)

            n_classes = logits_order_ambiguity.shape[-1]

        logits = torch.zeros((B, n_classes)).to(device)

        if order_ambiguity_mask.sum() < B:
            logits[~order_ambiguity_mask] = logits_no_order_ambiguity

        if order_ambiguity_mask.sum() > 0:
            logits[order_ambiguity_mask] = logits_order_ambiguity

        return logits, None
//  [python]
def spc(edge_index,args):
    set_random_seed(args.seed)
    if is_undirected(edge_index) == False:
        edge_index = to_undirected(edge_index)
    A=to_dense_adj(edge_index)[0]
    num_nodes=maybe_num_nodes(edge_index)
    w, v = torch.symeig(A,eigenvectors=True)
    sorted, indices = torch.sort(w)
    dim=args.embedding_dim
    indices=indices[-dim:]
    #indices=indices[:dim]
    v=v[:,indices]
    out=v.cpu().detach().clone()
    return out
//  [python]
def spc(edge_index,args):
    set_random_seed(args.seed)
    if is_undirected(edge_index) == False:
        edge_index = to_undirected(edge_index)
    A=to_dense_adj(edge_index)[0]
    num_nodes=maybe_num_nodes(edge_index)
    w, v = torch.linalg.eigh(A)
    sorted, indices = torch.sort(w)
    dim=args.embedding_dim
    indices=indices[-dim:]
    #indices=indices[:dim]
    v=v[:,indices]
    out=v.cpu().detach().clone()
    return out
//  [python]
def eigen_solve(A, mode):
    """solve for eigenpairs using a specified method"""
    if mode == 'arpack_eigsh':
        return arpack.eigsh(A, largest=True, tol=1e-4)
    elif mode == 'arpack_eigsh_mkl':
        return arpack.eigsh_mkl(A, largest=True, tol_dps=4)
    elif mode == 'torch_eigh':
        return torch.symeig(A, eigenvectors=True)
    elif mode == 'torch_lobpcg':
        return torch.lobpcg(A, k=1, largest=True, tol=1e-4)
    elif mode == 'scipy_eigsh':
        # For some reason scipy's eigsh requires slightly smaller tolerance
        # (1e-5 vs 1e-4) to reach equiavelent accuracy
        return splinalg.eigsh(A.numpy(), k=1, which="LA", tol=1e-5)
    elif mode == 'scipy_lobpcg':
        X = A.new_empty(A.size(0), 1).normal_()
        return splinalg.lobpcg(A.numpy(), X.numpy(), largest=True, tol=1e-4)
    else:
        raise ValueError
//  [python]
def eigen_solve(A, mode):
    """solve for eigenpairs using a specified method"""
    if mode == 'arpack_eigsh':
        return arpack.eigsh(A, largest=True, tol=1e-4)
    elif mode == 'arpack_eigsh_mkl':
        return arpack.eigsh_mkl(A, largest=True, tol_dps=4)
    elif mode == 'torch_eigh':
        return torch.linalg.eigh(A)
    elif mode == 'torch_lobpcg':
        return torch.lobpcg(A, k=1, largest=True, tol=1e-4)
    elif mode == 'scipy_eigsh':
        # For some reason scipy's eigsh requires slightly smaller tolerance
        # (1e-5 vs 1e-4) to reach equiavelent accuracy
        return splinalg.eigsh(A.numpy(), k=1, which="LA", tol=1e-5)
    elif mode == 'scipy_lobpcg':
        X = A.new_empty(A.size(0), 1).normal_()
        return splinalg.lobpcg(A.numpy(), X.numpy(), largest=True, tol=1e-4)
    else:
        raise ValueError
//  [python]
def norm(tensor: TensorType, dim: Optional[int] = None) -> torch.Tensor:
        """Tensor Norm"""
        return torch.norm(tensor, dim=dim)
//  [python]
def norm(tensor: TensorType, dim: Optional[int] = None) -> torch.Tensor:
        """Tensor Norm"""
        return torch.linalg.norm(tensor, dim=dim)
//  [python]
def eigen_decompose(
    laplacian: torch.Tensor, device: str
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Calculate eigenvalues and eigenvectors from the Laplacian matrix.

    Args:
        laplacian (torch.Tensor): Laplacian matrix
        device (str): Device to use for eigendecomposition.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tuple containing eigenvalues and eigenvectors.
    """
    laplacian = laplacian.float()

    if laplacian.device != device:
        laplacian = laplacian.to(device)

    lambdas, diffusion_map = torch.symeig(laplacian, eigenvectors=True)

    return lambdas, diffusion_map
//  [python]
def eigen_decompose(
    laplacian: torch.Tensor, device: str
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Calculate eigenvalues and eigenvectors from the Laplacian matrix.

    Args:
        laplacian (torch.Tensor): Laplacian matrix
        device (str): Device to use for eigendecomposition.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tuple containing eigenvalues and eigenvectors.
    """
    laplacian = laplacian.float()

    if laplacian.device != device:
        laplacian = laplacian.to(device)

    lambdas, diffusion_map = torch.linalg.eigh(laplacian)

    return lambdas, diffusion_map
//  [python]
def iterative_reweighted_ransac(points, pillars, w0, num_pillars, sigma2, stopping_delta=1e-2):
    """
    Args:
        points.bxyz [N, 4]: (batch, x, y, z)
        pillars: dict()
        num_pillars: number of pillars
        w0 [N, 1]: per-point weight
    Returns:
        points.plane_fitting_error [N]: dist from point to the final plane
    """
    assert len(w0.shape) == 2, "w0 must be a two dimensional tensor"
    w = w0
    point_xyz = points.bxyz[:, 1:4]
    for itr in range(50):
        center = segment_coo(
                     point_xyz*w,
                     points.pillar_idx,
                     dim_size=num_pillars,
                     reduce='sum'
                 )
        weight_sum = segment_coo(
                         w,
                         points.pillar_idx,
                         dim_size=num_pillars,
                         reduce='sum'
                     ) + 1e-6
        pillars.center = center / weight_sum

        point_d = point_xyz - pillars.center[points.pillar_idx]
        point_ddT = (w[:, :, None] * point_d[:, :, None]) * point_d[:, None, :]
        pillar_cov = segment_coo(
                         point_ddT,
                         points.pillar_idx,
                         dim_size=num_pillars,
                         reduce='mean'
                     )

        S, Q = torch.symeig(pillar_cov)
        pillars.normal = Q[:, :, 0]
        points.plane_fitting_error = (point_d * pillars.normal[points.pillar_idx]).sum(-1).abs()
        new_w = (sigma2 / (points.plane_fitting_error.square() + sigma2)).reshape(-1)
        dist_w = (0.5**2)/(point_d.square().sum(dim=-1) + 0.5**2)
        new_w = (new_w * dist_w).reshape(-1, 1)
        delta_w = (new_w - w).abs()
        if delta_w.max() < stopping_delta:
            break
        w = new_w

    return points, pillars
//  [python]
def iterative_reweighted_ransac(points, pillars, w0, num_pillars, sigma2, stopping_delta=1e-2):
    """
    Args:
        points.bxyz [N, 4]: (batch, x, y, z)
        pillars: dict()
        num_pillars: number of pillars
        w0 [N, 1]: per-point weight
    Returns:
        points.plane_fitting_error [N]: dist from point to the final plane
    """
    assert len(w0.shape) == 2, "w0 must be a two dimensional tensor"
    w = w0
    point_xyz = points.bxyz[:, 1:4]
    for itr in range(50):
        center = segment_coo(
                     point_xyz*w,
                     points.pillar_idx,
                     dim_size=num_pillars,
                     reduce='sum'
                 )
        weight_sum = segment_coo(
                         w,
                         points.pillar_idx,
                         dim_size=num_pillars,
                         reduce='sum'
                     ) + 1e-6
        pillars.center = center / weight_sum

        point_d = point_xyz - pillars.center[points.pillar_idx]
        point_ddT = (w[:, :, None] * point_d[:, :, None]) * point_d[:, None, :]
        pillar_cov = segment_coo(
                         point_ddT,
                         points.pillar_idx,
                         dim_size=num_pillars,
                         reduce='mean'
                     )

        S, Q = torch.linalg.eigh(pillar_cov)
        pillars.normal = Q[:, :, 0]
        points.plane_fitting_error = (point_d * pillars.normal[points.pillar_idx]).sum(-1).abs()
        new_w = (sigma2 / (points.plane_fitting_error.square() + sigma2)).reshape(-1)
        dist_w = (0.5**2)/(point_d.square().sum(dim=-1) + 0.5**2)
        new_w = (new_w * dist_w).reshape(-1, 1)
        delta_w = (new_w - w).abs()
        if delta_w.max() < stopping_delta:
            break
        w = new_w

    return points, pillars
//  [python]
def forward(self, data):
        """
        Args:
            data: (batch_size, num_samples, num_features) or (num_samples, num_features)
        Returns:
            output: (batch_size, num_samples, num_features) or (num_samples, num_features) in PCA domain.
        """
        n_dims = data.dim()

        assert n_dims in [2, 3], "data is expected 2D or 3D tensor."

        if n_dims == 2:
            data = data.unsqueeze(dim=0) # (batch_size, num_samples, num_features), where batch_size = 1.

        if self.training:
            if self.standardize:
                self.mean, self.std = torch.mean(data, dim=1), torch.std(data, dim=1)
                standardized = self.preprocess(data)
            else:
                standardized = data

            cov = torch.bmm(standardized.permute(0, 2, 1), standardized) / standardized.size(1)
            _, proj_matrix = torch.symeig(cov, eigenvectors=True) # Computed by assending order
            self.proj_matrix = torch.flip(proj_matrix, dims=(-1,))
        else:
            if self.proj_matrix is None:
                raise RuntimeError("proj_matrix is computed in advance.")

            if self.standardize:
                standardized = self.preprocess(data)
            else:
                standardized = data

        output = torch.bmm(standardized, self.proj_matrix)

        if n_dims == 2:
            output = output.squeeze(dim=0)

        return output