jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(p.data, alpha = self.weight_decay)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                self.d_a[m], self.Q_a[m] = torch.symeig(self.m_aa[m])
                self.d_g[m], self.Q_g[m] = torch.symeig(self.m_gg[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(p.data, alpha = self.weight_decay)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(self.m_aa[m])
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(self.m_gg[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        self.t += 0.1
        print(self.t)
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.symeig(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=False)
            V = V[:,indices]
            num_components = 1
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += self.t*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            #norm_direction = torch.norm(direction)
            #direction = direction/norm_direction

            '''
            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            stepSize = max(0.2,min(0.7,stepSize))
            #print(stepSize)
            '''
            stepSize = 1.0
            #stepSize = 0.1

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #new_draw = direction.clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        self.t += 0.1
        print(self.t)
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.linalg.eigh(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=False)
            V = V[:,indices]
            num_components = 1
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += self.t*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            #norm_direction = torch.norm(direction)
            #direction = direction/norm_direction

            '''
            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            stepSize = max(0.2,min(0.7,stepSize))
            #print(stepSize)
            '''
            stepSize = 1.0
            #stepSize = 0.1

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #new_draw = direction.clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def pca(self, X, n_components):
        """PCA linear proj from (N, din) to (N, dout)"""

        # Center Data
        X_mean = X - torch.mean(X, axis=0)

        # Compute Covariance Matrix
        cov_mat = torch.matmul(X_mean.transpose(0, 1), X_mean)

        # Compute Eigenvalues and Eigenvectors of the Covariance Matrix
        eigen_values, eigen_vectors = torch.symeig(cov_mat, eigenvectors=True)

        # Sort Eigenvectors in descending order
        sorted_index = torch.argsort(eigen_values, descending=True)
        sorted_eigenvectors = eigen_vectors[:, sorted_index]  # (D, D)

        # Select n components
        eigenvector_subset = sorted_eigenvectors[:, :n_components]  # (D, n)

        # Transform Data (B, D) -> (B, n)
        X_reduced = torch.matmul(X_mean, eigenvector_subset)

        return X_reduced
//  [python]
def pca(self, X, n_components):
        """PCA linear proj from (N, din) to (N, dout)"""

        # Center Data
        X_mean = X - torch.mean(X, axis=0)

        # Compute Covariance Matrix
        cov_mat = torch.matmul(X_mean.transpose(0, 1), X_mean)

        # Compute Eigenvalues and Eigenvectors of the Covariance Matrix
        eigen_values, eigen_vectors = torch.linalg.eigh(cov_mat)

        # Sort Eigenvectors in descending order
        sorted_index = torch.argsort(eigen_values, descending=True)
        sorted_eigenvectors = eigen_vectors[:, sorted_index]  # (D, D)

        # Select n components
        eigenvector_subset = sorted_eigenvectors[:, :n_components]  # (D, n)

        # Transform Data (B, D) -> (B, n)
        X_reduced = torch.matmul(X_mean, eigenvector_subset)

        return X_reduced
//  [python]
def test_eigh_2x2(dtype):

    # random PSD matrix
    A = torch.zeros((2, 10), dtype=dtype).uniform_()
    A = A @ torchiva.linalg.hermite(A)

    e1, v1 = torch.linalg.eigh(A)
    e2, v2 = torchiva.linalg.eigh_2x2(A)

    err_e = abs(e1 - e2).max()
    err_v = abs(abs(v1.T @ v2) - torch.eye(2)).max()

    assert err_e < tol
    assert err_v < tol
//  [python]
def measure_eigh_runtime(batch, dtype):
    rep = 20

    # random PSD matrix
    A = torch.zeros((batch, 2, 10), dtype=dtype).uniform_()
    A = A @ torchiva.linalg.hermite(A)

    A = A.to(0)

    e1, v1 = torch.linalg.eigh(A)
    ts = time.perf_counter()
    for i in range(rep):
        e1, v1 = torch.linalg.eigh(A)
    runtime_torch = (time.perf_counter() - ts) / rep

    e2, v2 = torchiva.linalg.eigh_2x2(A)
    ts = time.perf_counter()
    for i in range(rep):
        e2, v2 = torchiva.linalg.eigh_2x2(A)
    runtime_torchiva = (time.perf_counter() - ts) / rep


    print(f"batch={batch} dtype={dtype} torch={runtime_torch:.6f} torchiva={runtime_torchiva:.6f}")
//  [python]
def forward(ctx, K:torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        lambda_, U = torch.symeig(K)
        ctx.save_for_backward(lambda_, U)
        return lambda_, U
//  [python]
def forward(ctx, K:torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        lambda_, U = torch.linalg.eigh(K)
        ctx.save_for_backward(lambda_, U)
        return lambda_, U
//  [python]
def _compute_U_V_r(self, X, Y, ny_pts_x, ny_pts_y):
        penalty = self.penalty * X.shape[0]

        kx_nm = self.kernel(X, ny_pts_x)
        kx_mm = self.kernel(ny_pts_x, ny_pts_x)

        ky_nm = self.kernel(Y, ny_pts_y)
        ky_mm = self.kernel(ny_pts_y, ny_pts_y)

        toinv = kx_nm.T @ kx_nm + penalty * kx_mm
        edec1 = torch.symeig(toinv)
        r = kx_mm.shape[0]
        edec1_topr = fetch_topr_eigensystem(edec1, r=r)  # This is for sorting
        inv_edec1 = pseudo_inverse(edec1_topr, EIGENVALUE_EPS[edec1_topr.eigenvectors.dtype])
        V = inv_edec1.eigenvectors * inv_edec1.eigenvalues
        U = ky_nm.T @ kx_nm @ inv_edec1.eigenvectors
        U = torch.linalg.solve(id_regularize(ky_mm, 1e-8), U)
        return U, V, inv_edec1.num_nonzero or r
//  [python]
def _compute_U_V_r(self, X, Y, ny_pts_x, ny_pts_y):
        penalty = self.penalty * X.shape[0]

        kx_nm = self.kernel(X, ny_pts_x)
        kx_mm = self.kernel(ny_pts_x, ny_pts_x)

        ky_nm = self.kernel(Y, ny_pts_y)
        ky_mm = self.kernel(ny_pts_y, ny_pts_y)

        toinv = kx_nm.T @ kx_nm + penalty * kx_mm
        edec1 = torch.linalg.eigh(toinv)
        r = kx_mm.shape[0]
        edec1_topr = fetch_topr_eigensystem(edec1, r=r)  # This is for sorting
        inv_edec1 = pseudo_inverse(edec1_topr, EIGENVALUE_EPS[edec1_topr.eigenvectors.dtype])
        V = inv_edec1.eigenvectors * inv_edec1.eigenvalues
        U = ky_nm.T @ kx_nm @ inv_edec1.eigenvectors
        U = torch.linalg.solve(id_regularize(ky_mm, 1e-8), U)
        return U, V, inv_edec1.num_nonzero or r
//  [python]
def forward(ctx, A):
        """Calculate the eigenvalues and eigenvectors of a symmetric matrix.

        Parameters
        ----------
        A: torch.Tensor
            The symmetric matrix to be diagonalized. Shape: ``(*BA, q, q)``.

        Returns
        -------
        eival: torch.Tensor
            Eigenvalues of the linear operator.
        eivec: torch.Tensor
            Eigenvectors of the linear operator.

        """
        eival, eivec = torch.symeig(A)
        ctx.save_for_backward(eival, eivec)
        return eival, eivec
//  [python]
def forward(ctx, A):
        """Calculate the eigenvalues and eigenvectors of a symmetric matrix.

        Parameters
        ----------
        A: torch.Tensor
            The symmetric matrix to be diagonalized. Shape: ``(*BA, q, q)``.

        Returns
        -------
        eival: torch.Tensor
            Eigenvalues of the linear operator.
        eivec: torch.Tensor
            Eigenvectors of the linear operator.

        """
        eival, eivec = torch.linalg.eigh(A)
        ctx.save_for_backward(eival, eivec)
        return eival, eivec
//  [python]
def quaternion_average(self, Q1, weights):
        Q1 = Q1 + eps
        # print("\n\n-------Quaternions-----------\n{}".format(Q1))
        weights = F.normalize(weights, p=1, dim=1)
        outer = weights[(...,) + (None,) * 2] * (Q1.unsqueeze(dim=-1) @ Q1.unsqueeze(dim=3))
        M = outer.sum(1, keepdim=True)
        M = M + (1e-6) * torch.randn_like(M)

        eigs, eigvect = torch.symeig(M.view(-1, M.shape[-2], M.shape[-1]).cpu(), eigenvectors=True)
        eigs = eigs.to(engine.device)
        eigvect = eigvect.to(engine.device)
        normalized_eigs = torch.div(eigs, eigs.sum(dim=1).unsqueeze(1))
        normalized_max_eig = normalized_eigs[:, -1].view(M.shape[0], M.shape[-3], 1)
        avg_quats = eigvect[:, -1, :].view(M.shape[:-1])

        return normalized_max_eig, avg_quats
//  [python]
def quaternion_average(self, Q1, weights):
        Q1 = Q1 + eps
        # print("\n\n-------Quaternions-----------\n{}".format(Q1))
        weights = F.normalize(weights, p=1, dim=1)
        outer = weights[(...,) + (None,) * 2] * (Q1.unsqueeze(dim=-1) @ Q1.unsqueeze(dim=3))
        M = outer.sum(1, keepdim=True)
        M = M + (1e-6) * torch.randn_like(M)

        eigs, eigvect = torch.linalg.eigh(M.view(-1, M.shape[-2], M.shape[-1]).cpu())
        eigs = eigs.to(engine.device)
        eigvect = eigvect.to(engine.device)
        normalized_eigs = torch.div(eigs, eigs.sum(dim=1).unsqueeze(1))
        normalized_max_eig = normalized_eigs[:, -1].view(M.shape[0], M.shape[-3], 1)
        avg_quats = eigvect[:, -1, :].view(M.shape[:-1])

        return normalized_max_eig, avg_quats
//  [python]
def eigh(data, meta=None, sizes=(1, 1), order_by_magnitude=False, ad_decomp_reg=1.0e-12):
    real_dtype= data.real.dtype if data.is_complex() else data.dtype
    Sdata = torch.zeros((sizes[0],), dtype=real_dtype, device=data.device)
    Udata = torch.zeros((sizes[1],), dtype=data.dtype, device=data.device)
    if meta is not None:
        if order_by_magnitude:
            reg = torch.as_tensor(ad_decomp_reg, dtype=real_dtype, device=data.device)
            f = lambda x: SYMEIG.apply(x, reg)
        else:
            f = lambda x: torch.symeig(x)
        for (sl, D, slU, DU, slS) in meta:
            S, U = f(data[slice(*sl)].view(D))
            Sdata[slice(*slS)] = S
            Udata[slice(*slU)].view(DU)[:] = U
        return Sdata, Udata
    return torch.symeig(data)
//  [python]
def eigh(data, meta=None, sizes=(1, 1), order_by_magnitude=False, ad_decomp_reg=1.0e-12):
    real_dtype= data.real.dtype if data.is_complex() else data.dtype
    Sdata = torch.zeros((sizes[0],), dtype=real_dtype, device=data.device)
    Udata = torch.zeros((sizes[1],), dtype=data.dtype, device=data.device)
    if meta is not None:
        if order_by_magnitude:
            reg = torch.as_tensor(ad_decomp_reg, dtype=real_dtype, device=data.device)
            f = lambda x: SYMEIG.apply(x, reg)
        else:
            f = lambda x: torch.linalg.eigh(x)
        for (sl, D, slU, DU, slS) in meta:
            S, U = f(data[slice(*sl)].view(D))
            Sdata[slice(*slS)] = S
            Udata[slice(*slU)].view(DU)[:] = U
        return Sdata, Udata
    return torch.linalg.eigh(data)
//  [python]
def pca(image):
    """
    Compute the major components of an image. The Image is treated as a
    distribution.
    Parameters
    ----------
    image: Image or 2DArray (N, M)
            Image to be used as distribution
    Returns
    -------
    cog_x: Skalar
            X-position of the distributions center of gravity
    cog_y: Skalar
            Y-position of the distributions center of gravity
    psi: Skalar
            Angle between first mjor component and x-axis
    """
    torch.set_printoptions(precision=16)

    pix_x, pix_y, image = im_to_array_value(image)

    cog_x = (torch.sum(pix_x * image, axis=1) / torch.sum(image, axis=1)).unsqueeze(-1)
    cog_y = (torch.sum(pix_y * image, axis=1) / torch.sum(image, axis=1)).unsqueeze(-1)

    delta_x = pix_x - cog_x
    delta_y = pix_y - cog_y

    inp = torch.cat([delta_x.unsqueeze(1), delta_y.unsqueeze(1)], dim=1)

    cov_w = bmul(
        (cog_x - 1 * torch.sum(image * image, axis=1).unsqueeze(-1) / cog_x).squeeze(1),
        (torch.matmul(image.unsqueeze(1) * inp, inp.transpose(1, 2))),
    )

    eig_vals_torch, eig_vecs_torch = torch.symeig(cov_w, eigenvectors=True)

    psi_torch = torch.atan(eig_vecs_torch[:, 1, 1] / eig_vecs_torch[:, 0, 1])

    return cog_x, cog_y, psi_torch
//  [python]
def pca(image):
    """
    Compute the major components of an image. The Image is treated as a
    distribution.
    Parameters
    ----------
    image: Image or 2DArray (N, M)
            Image to be used as distribution
    Returns
    -------
    cog_x: Skalar
            X-position of the distributions center of gravity
    cog_y: Skalar
            Y-position of the distributions center of gravity
    psi: Skalar
            Angle between first mjor component and x-axis
    """
    torch.set_printoptions(precision=16)

    pix_x, pix_y, image = im_to_array_value(image)

    cog_x = (torch.sum(pix_x * image, axis=1) / torch.sum(image, axis=1)).unsqueeze(-1)
    cog_y = (torch.sum(pix_y * image, axis=1) / torch.sum(image, axis=1)).unsqueeze(-1)

    delta_x = pix_x - cog_x
    delta_y = pix_y - cog_y

    inp = torch.cat([delta_x.unsqueeze(1), delta_y.unsqueeze(1)], dim=1)

    cov_w = bmul(
        (cog_x - 1 * torch.sum(image * image, axis=1).unsqueeze(-1) / cog_x).squeeze(1),
        (torch.matmul(image.unsqueeze(1) * inp, inp.transpose(1, 2))),
    )

    eig_vals_torch, eig_vecs_torch = torch.linalg.eigh(cov_w, UPLO="U")

    psi_torch = torch.atan(eig_vecs_torch[:, 1, 1] / eig_vecs_torch[:, 0, 1])

    return cog_x, cog_y, psi_torch
//  [python]
def step(self, closure):
        """Solves a subproblem.
        Arguments:
            closure (callable): a closure that reevaluates the model and returns the loss.
        """
        closure = torch.enable_grad()(closure)

        for group in self.param_groups:
            params = group['params']

            L = group['L']
            subsolver = group['subsolver']
            max_iters = group['max_iters']
            subsolver_args = group['subsolver_args']
            max_iters_outer = group['max_iters_outer']

            df = tuple_to_vec.tuple_to_vector(
                torch.autograd.grad(closure(), list(params), create_graph=True))

            v_flat = torch.zeros_like(df)
            v = tuple_to_vec.rollup_vector(v_flat, list(params))

            if subsolver is None:
                full_hessian = derivatives.flat_hessian(df, list(params)).to(torch.double)
                eigenvalues, eigenvectors = torch.symeig(full_hessian)

            for _ in range(max_iters_outer):
                D3xx, Hx = derivatives.third_derivative_vec(
                    closure, list(params), v, flat=True)
                with torch.no_grad():
                    D3xx = D3xx.to(torch.double)
                    Hx = Hx.to(torch.double)
                    Lv3 = v_flat * (v_flat.norm() ** 2 * L)
                    g = df.add(Hx).add(D3xx, alpha=0.5).add(Lv3)

                if self._check_stopping_condition(closure, params, v, g.norm()):
                    self._add_v(params, v)
                    return True

                with torch.no_grad():
                    c = g.div(2. + math.sqrt(2)).sub(Hx).sub(Lv3)

                if subsolver is None:
                    v_flat = exact(L, c.detach(), T=eigenvalues, U=eigenvectors)
                else:
                    v_flat = iterative(params, closure, L, c.detach(),
                                  subsolver, subsolver_args, max_iters)
                with torch.no_grad():
                    v = tuple_to_vec.rollup_vector(v_flat, list(params))

            self._add_v(params, v)
        return False
//  [python]
def step(self, closure):
        """Solves a subproblem.
        Arguments:
            closure (callable): a closure that reevaluates the model and returns the loss.
        """
        closure = torch.enable_grad()(closure)

        for group in self.param_groups:
            params = group['params']

            L = group['L']
            subsolver = group['subsolver']
            max_iters = group['max_iters']
            subsolver_args = group['subsolver_args']
            max_iters_outer = group['max_iters_outer']

            df = tuple_to_vec.tuple_to_vector(
                torch.autograd.grad(closure(), list(params), create_graph=True))

            v_flat = torch.zeros_like(df)
            v = tuple_to_vec.rollup_vector(v_flat, list(params))

            if subsolver is None:
                full_hessian = derivatives.flat_hessian(df, list(params)).to(torch.double)
                eigenvalues, eigenvectors = torch.linalg.eigh(full_hessian)

            for _ in range(max_iters_outer):
                D3xx, Hx = derivatives.third_derivative_vec(
                    closure, list(params), v, flat=True)
                with torch.no_grad():
                    D3xx = D3xx.to(torch.double)
                    Hx = Hx.to(torch.double)
                    Lv3 = v_flat * (v_flat.norm() ** 2 * L)
                    g = df.add(Hx).add(D3xx, alpha=0.5).add(Lv3)

                if self._check_stopping_condition(closure, params, v, g.norm()):
                    self._add_v(params, v)
                    return True

                with torch.no_grad():
                    c = g.div(2. + math.sqrt(2)).sub(Hx).sub(Lv3)

                if subsolver is None:
                    v_flat = exact(L, c.detach(), T=eigenvalues, U=eigenvectors)
                else:
                    v_flat = iterative(params, closure, L, c.detach(),
                                  subsolver, subsolver_args, max_iters)
                with torch.no_grad():
                    v = tuple_to_vec.rollup_vector(v_flat, list(params))

            self._add_v(params, v)
        return False
//  [python]
def pca_calc(
    X: list[torch.Tensor], ignore_masks: list[torch.Tensor] | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Run PCA on a list of batched data. Returns the eigenvalues and eigenvectors.
    """
    # Run GC and cleanup GPU memory
    cleanup_memory()

    H = None
    for idx, X_batch in enumerate(X):
        if ignore_masks:
            X_batch[ignore_masks[idx] == 0] = 0

        X_batch = X_batch.double().to(device=config.device)
        H_batch = torch.sum(X_batch.mT @ X_batch, dim=0)  # sum over the batch dimension.
        H = H_batch if H is None else H + H_batch

    damp = 0.01 * torch.mean(torch.diag(H))
    diag = torch.arange(H.shape[-1]).to(device=config.device)
    H[diag, diag] = H[diag, diag] + damp
    X_eig = torch.symeig(H, eigenvectors=True)
    del H
    index = torch.argsort(X_eig[0], descending=True)
    eig_val = X_eig[0][index]
    eigen_vec = X_eig[1][:, index]
    return eig_val, eigen_vec
//  [python]
def pca_calc(
    X: list[torch.Tensor], ignore_masks: list[torch.Tensor] | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Run PCA on a list of batched data. Returns the eigenvalues and eigenvectors.
    """
    # Run GC and cleanup GPU memory
    cleanup_memory()

    H = None
    for idx, X_batch in enumerate(X):
        if ignore_masks:
            X_batch[ignore_masks[idx] == 0] = 0

        X_batch = X_batch.double().to(device=config.device)
        H_batch = torch.sum(X_batch.mT @ X_batch, dim=0)  # sum over the batch dimension.
        H = H_batch if H is None else H + H_batch

    damp = 0.01 * torch.mean(torch.diag(H))
    diag = torch.arange(H.shape[-1]).to(device=config.device)
    H[diag, diag] = H[diag, diag] + damp
    X_eig = torch.linalg.eigh(H)
    del H
    index = torch.argsort(X_eig[0], descending=True)
    eig_val = X_eig[0][index]
    eigen_vec = X_eig[1][:, index]
    return eig_val, eigen_vec
//  [python]
def forward(self, S):
        """Forward Pass.
        Args:
            S (:obj:`torch.Tensor`): input to be forwarded. (N_sample, Npx)
        Returns:
            :obj:`torch.Tensor`: output: (N_sample, Npx)
        """
        N_sample, N_px = S.shape[0], self.tau.shape[0]
        y = torch.zeros((N_sample, N_px))
        for i in range(N_sample): # Loop to handle linalg.eigh: broadcasting can be slower
            Ds, Vs = torch.symeig(S[i]) # (Nch, Nch), (Nch, Nch)
            idx = Ds > 0  # To avoid np.sqrt() issues.
            Ds, Vs = Ds[idx], Vs[:, idx]
            y[i] = torch.linalg.norm(self.D.conj().T @ (Vs * torch.sqrt(Ds)), axis=1) ** 2 # (Npx, Nch) dot ((Nch, Nch) * (Nch, Nch))
        y -= self.tau
        return y
//  [python]
def forward(self, S):
        """Forward Pass.
        Args:
            S (:obj:`torch.Tensor`): input to be forwarded. (N_sample, Npx)
        Returns:
            :obj:`torch.Tensor`: output: (N_sample, Npx)
        """
        N_sample, N_px = S.shape[0], self.tau.shape[0]
        y = torch.zeros((N_sample, N_px))
        for i in range(N_sample): # Loop to handle linalg.eigh: broadcasting can be slower
            Ds, Vs = torch.linalg.eigh(S[i]) # (Nch, Nch), (Nch, Nch)
            idx = Ds > 0  # To avoid np.sqrt() issues.
            Ds, Vs = Ds[idx], Vs[:, idx]
            y[i] = torch.linalg.norm(self.D.conj().T @ (Vs * torch.sqrt(Ds)), axis=1) ** 2 # (Npx, Nch) dot ((Nch, Nch) * (Nch, Nch))
        y -= self.tau
        return y
//  [python]
def precondition(self):
        hessian = self.model.hessian.data
        grads = [p.grad for p in self.model.parameters() if p.grad is not None]
        if self.config.absolute:
            L, Q = torch.symeig(hessian)
            hessian = Q @ torch.abs(torch.diag(L)) @ Q.T
        diag = torch.diagonal(hessian)
        diag += self.config.damping
        g = parameters_to_vector(grads)
        vector_to_parameters(torch.linalg.solve(hessian, g), grads)
//  [python]
def precondition(self):
        hessian = self.model.hessian.data
        grads = [p.grad for p in self.model.parameters() if p.grad is not None]
        if self.config.absolute:
            L, Q = torch.linalg.eigh(hessian)
            hessian = Q @ torch.abs(torch.diag(L)) @ Q.T
        diag = torch.diagonal(hessian)
        diag += self.config.damping
        g = parameters_to_vector(grads)
        vector_to_parameters(torch.linalg.solve(hessian, g), grads)
//  [python]
def test_sqrt_inv_matmul_no_lhs(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = linear_op.clone().detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        sqrt_inv_matmul_res = linear_op.sqrt_inv_matmul(rhs)
        evals, evecs = torch.symeig(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = matrix_inv_root @ rhs_copy

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def test_sqrt_inv_matmul_no_lhs(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = linear_op.clone().detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        sqrt_inv_matmul_res = linear_op.sqrt_inv_matmul(rhs)
        evals, evecs = torch.linalg.eigh(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = matrix_inv_root @ rhs_copy

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def test_sqrt_inv_matmul(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = linear_op.clone().detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        lhs = torch.randn(*linear_op.shape[:-2], 2, linear_op.size(-1)).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        lhs_copy = lhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        sqrt_inv_matmul_res, inv_quad_res = linear_op.sqrt_inv_matmul(rhs, lhs)
        evals, evecs = torch.symeig(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = lhs_copy @ matrix_inv_root @ rhs_copy
        inv_quad_actual = (lhs_copy @ matrix_inv_root).pow(2).sum(dim=-1)

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(inv_quad_res, inv_quad_actual, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def test_sqrt_inv_matmul(self):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        if len(linear_op.batch_shape):
            return

        linear_op_copy = linear_op.clone().detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)

        # Create a test right hand side and left hand side
        rhs = torch.randn(*linear_op.shape[:-1], 3).requires_grad_(True)
        lhs = torch.randn(*linear_op.shape[:-2], 2, linear_op.size(-1)).requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        lhs_copy = lhs.clone().detach().requires_grad_(True)

        # Perform forward pass
        sqrt_inv_matmul_res, inv_quad_res = linear_op.sqrt_inv_matmul(rhs, lhs)
        evals, evecs = torch.linalg.eigh(evaluated)
        matrix_inv_root = evecs @ (evals.sqrt().reciprocal().unsqueeze(-1) * evecs.mT)
        sqrt_inv_matmul_actual = lhs_copy @ matrix_inv_root @ rhs_copy
        inv_quad_actual = (lhs_copy @ matrix_inv_root).pow(2).sum(dim=-1)

        # Check forward pass
        self.assertAllClose(sqrt_inv_matmul_res, sqrt_inv_matmul_actual, **self.tolerances["sqrt_inv_matmul"])
        self.assertAllClose(inv_quad_res, inv_quad_actual, **self.tolerances["sqrt_inv_matmul"])
//  [python]
def load_norm_label(self, norm_path, H, W, depth, K):
        depth = torch.from_numpy(depth).squeeze()
        K = torch.Tensor([[K[0], 0 ,K[2]], 
                      [0, K[1], K[3]], 
                      [0, 0, 1]])
        K_inv = K.inverse()

        y, x = torch.meshgrid([torch.arange(0, 480, dtype=torch.float32),
                            torch.arange(0, 640, dtype=torch.float32)], indexing='ij')
        x = x.reshape(1, 480*640)
        y = y.reshape(1, 480*640)
        ones = torch.ones_like(x)
        coord_2d = torch.cat((x, y, ones), dim=0)

        coord_3d = torch.matmul(K_inv, coord_2d).view(3, 480, 640)
        coord_3d = (coord_3d * depth[None, :])[None, :]
        coord_3d_mean = self.avg(coord_3d)

        uf_coord_3d = self.unfold(coord_3d.permute(1, 0, 2, 3))
        coord_3d_decenter = uf_coord_3d - coord_3d_mean.view(3, 1, (480-6)*(640-6))
        coord_3d_decenter = coord_3d_decenter.permute(2, 0, 1)
        cov = torch.bmm(coord_3d_decenter, coord_3d_decenter.permute(0, 2, 1))
        
        eig = torch.symeig(cov)
        #svd = torch.linalg.svd(coord_3d_decenter)
        normal = (eig.eigenvectors)[:, :, 0].float()
        #normal = (svd[1])[:, 2, :]
        normal = self.pad(normal.permute(1, 0).view(1, 3, (480-6), (640-6)))
        
        orient_mask = (torch.sum(normal * coord_3d, axis=1) < 0).unsqueeze(1)
        normal = normal * orient_mask - normal * (~orient_mask)
        gt_normal = normal.squeeze().permute(1, 2, 0).numpy()
        return gt_normal
//  [python]
def load_norm_label(self, norm_path, H, W, depth, K):
        depth = torch.from_numpy(depth).squeeze()
        K = torch.Tensor([[K[0], 0 ,K[2]], 
                      [0, K[1], K[3]], 
                      [0, 0, 1]])
        K_inv = K.inverse()

        y, x = torch.meshgrid([torch.arange(0, 480, dtype=torch.float32),
                            torch.arange(0, 640, dtype=torch.float32)], indexing='ij')
        x = x.reshape(1, 480*640)
        y = y.reshape(1, 480*640)
        ones = torch.ones_like(x)
        coord_2d = torch.cat((x, y, ones), dim=0)

        coord_3d = torch.matmul(K_inv, coord_2d).view(3, 480, 640)
        coord_3d = (coord_3d * depth[None, :])[None, :]
        coord_3d_mean = self.avg(coord_3d)

        uf_coord_3d = self.unfold(coord_3d.permute(1, 0, 2, 3))
        coord_3d_decenter = uf_coord_3d - coord_3d_mean.view(3, 1, (480-6)*(640-6))
        coord_3d_decenter = coord_3d_decenter.permute(2, 0, 1)
        cov = torch.bmm(coord_3d_decenter, coord_3d_decenter.permute(0, 2, 1))
        
        eig = torch.linalg.eigh(cov)
        #svd = torch.linalg.svd(coord_3d_decenter)
        normal = (eig[1])[:, :, 0].float()
        #normal = (svd[1])[:, 2, :]
        normal = self.pad(normal.permute(1, 0).view(1, 3, (480-6), (640-6)))
        
        orient_mask = (torch.sum(normal * coord_3d, axis=1) < 0).unsqueeze(1)
        normal = normal * orient_mask - normal * (~orient_mask)
        gt_normal = normal.squeeze().permute(1, 2, 0).numpy()
        return gt_normal
//  [python]
def compute_cov_mat(patches,rank):

    import time

    with torch.no_grad():
        # -- cov mat --
        bsize,num,pdim = patches.shape
        covMat = torch.matmul(patches.transpose(2,1),patches)
        covMat /= num

        # -- eigen stuff --

        # start = time.perf_counter()
        # eigVals,eigVecs = torch.linalg.eig(covMat)
        # end = time.perf_counter() - start

        # print("Eig Lin Time: ",end)
        # print("a: ",eigVals.shape,eigVecs.shape)

        # eigVals = torch.real(eigVals)#[...,:rank])
        # eigVecs = torch.real(eigVecs[...,:rank])
        # eigVals[...,rank:] = 0
        # print(eigVals[0])
        # print(eigVecs[0,0,:])

        # -- svd stuff --
        eigVals,eigVecs = torch.symeig(covMat,eigenvectors=True)
        # print(eigVals[0])
        eigVals= torch.flip(eigVals,dims=(1,))
        # print(eigVals[0])
        # print(eigVecs[0,0,:])
        eigVecs = torch.flip(eigVecs,dims=(2,))[...,:rank]
        # print(eigVals[0])
        # print(eigVecs[0,0,:])

        # print("eigh: ",end)
        # print("b: ",eigVals.shape,eigVecs.shape,bsize)

        # -- qr stuff --
        # start = time.perf_counter()
        # Q,R = torch.linalg.qr(covMat)
        # end = time.perf_counter() - start
        # print("qr: ",end)
        # print("QR: ",Q.shape,R.shape)

        # -- chol stuff --
        # start = time.perf_counter()
        # L,info = torch.linalg.cholesky_ex(covMat)
        # end = time.perf_counter() - start
        # print("chol: ",end)
        # print("chol: ",L.shape)

        # -- eig --
        # eigVals = rearrange(eigVals,'(b c) p -> b c p',b=bsize)
        # eigVecs = rearrange(eigVecs,'(b c) p p2 -> b c p p2',b=bsize)

        # -- rank --

    return covMat,eigVals,eigVecs
//  [python]
def compute_cov_mat(patches,rank):

    import time

    with torch.no_grad():
        # -- cov mat --
        bsize,num,pdim = patches.shape
        covMat = torch.matmul(patches.transpose(2,1),patches)
        covMat /= num

        # -- eigen stuff --

        # start = time.perf_counter()
        # eigVals,eigVecs = torch.linalg.eig(covMat)
        # end = time.perf_counter() - start

        # print("Eig Lin Time: ",end)
        # print("a: ",eigVals.shape,eigVecs.shape)

        # eigVals = torch.real(eigVals)#[...,:rank])
        # eigVecs = torch.real(eigVecs[...,:rank])
        # eigVals[...,rank:] = 0
        # print(eigVals[0])
        # print(eigVecs[0,0,:])

        # -- svd stuff --
        eigVals,eigVecs = torch.linalg.eigh(covMat)
        # print(eigVals[0])
        eigVals= torch.flip(eigVals,dims=(1,))
        # print(eigVals[0])
        # print(eigVecs[0,0,:])
        eigVecs = torch.flip(eigVecs,dims=(2,))[...,:rank]
        # print(eigVals[0])
        # print(eigVecs[0,0,:])

        # print("eigh: ",end)
        # print("b: ",eigVals.shape,eigVecs.shape,bsize)

        # -- qr stuff --
        # start = time.perf_counter()
        # Q,R = torch.linalg.qr(covMat)
        # end = time.perf_counter() - start
        # print("qr: ",end)
        # print("QR: ",Q.shape,R.shape)

        # -- chol stuff --
        # start = time.perf_counter()
        # L,info = torch.linalg.cholesky_ex(covMat)
        # end = time.perf_counter() - start
        # print("chol: ",end)
        # print("chol: ",L.shape)

        # -- eig --
        # eigVals = rearrange(eigVals,'(b c) p -> b c p',b=bsize)
        # eigVecs = rearrange(eigVecs,'(b c) p p2 -> b c p p2',b=bsize)

        # -- rank --

    return covMat,eigVals,eigVecs
//  [python]
def random_walk_covariance_static(
    gossip_matrix: torch.Tensor, gamma: float, start_at: int = 1
) -> torch.Tensor:
    """
    Asymptotic covariance `E[x x^T]` in the random walk process
    `x <- W @ (sqrt(gamma) x + n)`,
    where `x` is a vector containing one scalar per worker
    and `n` is i.i.d. standard normal noise.
    """
    W = gossip_matrix
    w_is_symmetric = W.allclose(W.T)
    if w_is_symmetric:
        L, Q = torch.symeig(W)
        numerator = L.square() if start_at == 1 else 1
        diag = numerator / (1 - gamma * L.square())
        return (Q * diag) @ Q.T  # = Q @ torch.diag(diag) @ Q.T
    else:
        rhs = W @ W.T if start_at == 1 else torch.eye(len(W))
        return solve_discrete_lyapunov(sqrt(gamma) * W, rhs)
//  [python]
def random_walk_covariance_static(
    gossip_matrix: torch.Tensor, gamma: float, start_at: int = 1
) -> torch.Tensor:
    """
    Asymptotic covariance `E[x x^T]` in the random walk process
    `x <- W @ (sqrt(gamma) x + n)`,
    where `x` is a vector containing one scalar per worker
    and `n` is i.i.d. standard normal noise.
    """
    W = gossip_matrix
    w_is_symmetric = W.allclose(W.T)
    if w_is_symmetric:
        L, Q = torch.linalg.eigh(W)
        numerator = L.square() if start_at == 1 else 1
        diag = numerator / (1 - gamma * L.square())
        return (Q * diag) @ Q.T  # = Q @ torch.diag(diag) @ Q.T
    else:
        rhs = W @ W.T if start_at == 1 else torch.eye(len(W))
        return solve_discrete_lyapunov(sqrt(gamma) * W, rhs)
//  [python]
def __eigs(self, C):
        """
        get eigenvalues and eigenvectors, sort them.
        Parameters
        ----------
        C: torch.Tensor
            The RGB values of skin-colored pixels.
        Returns
        -------
        Λ: torch.Tensor
            The eigenvalues of the correlation matrix
        U: torch.Tensor
            The (sorted) eigenvectors of the correlation matrix
        """
        # get eigenvectors and sort them according to eigenvalues (largest first)
        L, U = torch.symeig(C, eigenvectors=True)  # dim Λ: 3 | dim U: 3x3
        idx = torch.argsort(L)  # dim: 3x1
        idx = idx.flip(0)  # dim: 1x3
        L_ = L[idx]  # dim: 3
        U_ = U[:, idx]  # dim: 3x3

        return L_, U_
//  [python]
def __eigs(self, C):
        """
        get eigenvalues and eigenvectors, sort them.
        Parameters
        ----------
        C: torch.Tensor
            The RGB values of skin-colored pixels.
        Returns
        -------
        Λ: torch.Tensor
            The eigenvalues of the correlation matrix
        U: torch.Tensor
            The (sorted) eigenvectors of the correlation matrix
        """
        # get eigenvectors and sort them according to eigenvalues (largest first)
        L, U = torch.linalg.eigh(C)  # dim Λ: 3 | dim U: 3x3
        idx = torch.argsort(L)  # dim: 3x1
        idx = idx.flip(0)  # dim: 1x3
        L_ = L[idx]  # dim: 3
        U_ = U[:, idx]  # dim: 3x3

        return L_, U_
//  [python]
def test_project():
    """ Given a cMPS, check that the cMPS is unchanged after projecting a 
        unitary gate U without altering its bond dimension
    """
    dtype=torch.float64
    device='cpu'

    Q = torch.rand(8,8, dtype=dtype, device=device)
    R = torch.rand(4,8,8, dtype=dtype, device=device)
    beta = 10*torch.rand(1, dtype=dtype, device=device).item()
    mps = cmps(Q, R)

    A = torch.rand(8,8, dtype=dtype, device=device)
    A = A + A.t()
    w, U = torch.symeig(A)
    mps1 = mps.project(U)
   
    assert np.isclose(Fidelity(mps, mps1, beta), 0.5*ln_ovlp(mps, mps, beta))
//  [python]
def test_project():
    """ Given a cMPS, check that the cMPS is unchanged after projecting a 
        unitary gate U without altering its bond dimension
    """
    dtype=torch.float64
    device='cpu'

    Q = torch.rand(8,8, dtype=dtype, device=device)
    R = torch.rand(4,8,8, dtype=dtype, device=device)
    beta = 10*torch.rand(1, dtype=dtype, device=device).item()
    mps = cmps(Q, R)

    A = torch.rand(8,8, dtype=dtype, device=device)
    A = A + A.t()
    _, U = torch.linalg.eigh(A)
    mps1 = mps.project(U)
   
    assert np.isclose(Fidelity(mps, mps1, beta), 0.5*ln_ovlp(mps, mps, beta))
//  [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.symeig(voxel_ddT, eigenvectors=True)
                    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