jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype,
                                                           device=device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.linalg.solve(A, b)
        x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy())).to(dtype=dtype, device=device)
        self.assertEqual(x.data, x_exp)
//  [python]
def __init__(self, nfunc, edging=False, implicit=True, eps=-1.5, neddvalgrad=False, calcsd=0, *argin, **kwargin):
        nn.Module.__init__(self)
        self.implicit = implicit
        self.edging = edging
        self.logsigmoid = torch.nn.LogSigmoid()
        self.kf = MyKF(nfunc, *argin, **kwargin)
        self.softdot = functions.softdot1.apply if neddvalgrad else functions.softdot.apply
        self.neddvalgrad = neddvalgrad
        self.calcsd = calcsd
        self.lusolve = torch.linalg.solve if hasattr(torch,'linalg') and hasattr(torch.linalg,'solve') else lambda C,b: torch.solve(b,C)[0]
        self.eps = Parameter(eps*torch.ones(nfunc))
//  [python]
def __init__(self, nfunc, edging=False, implicit=True, eps=-1.5, neddvalgrad=False, calcsd=0, *argin, **kwargin):
        nn.Module.__init__(self)
        self.implicit = implicit
        self.edging = edging
        self.logsigmoid = torch.nn.LogSigmoid()
        self.kf = MyKF(nfunc, *argin, **kwargin)
        self.softdot = functions.softdot1.apply if neddvalgrad else functions.softdot.apply
        self.neddvalgrad = neddvalgrad
        self.calcsd = calcsd
        self.lusolve = torch.linalg.solve
        self.eps = Parameter(eps*torch.ones(nfunc))
//  [python]
def unproject2_aligned_depth(p2d, Kdist, Mdist, distCoeffs):
    pn2d = torch.linalg.solve(Kdist, torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0])).cuda().double()), axis=0)).T
    k = torch.cat((distCoeffs[:5], torch.zeros(12-5).cuda().double()), axis=0)
    x0 = pn2d[:,0]
    y0 = pn2d[:,1]
    x = x0
    y = y0
    for iter in range(5):
        r2 = x*x + y*y
        icdist = (1 + ((k[7]*r2 + k[6])*r2 + k[5])*r2)/(1 + ((k[4]*r2 + k[1])*r2 + k[0])*r2)
        deltaX = 2*k[2]*x*y + k[3]*(r2 + 2*x*x)+ k[8]*r2+k[9]*r2*r2
        deltaY = k[2]*(r2 + 2*y*y) + 2*k[3]*x*y+ k[10]*r2+k[11]*r2*r2
        x = (x0 - deltaX)*icdist
        y = (y0 - deltaY)*icdist
    
    pn2d = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), axis=1)
    
    p3d = torch.cat([pn2d * (p2d[:,2:3]*0.001), p2d[:,2:3]*0.001], axis=1)
    p3d = torch.matmul(torch.inverse(torch.cat([Mdist, torch.DoubleTensor([[0,0,0,1]]).cuda()], dim=0)), torch.cat((p3d, torch.ones((p3d.shape[0], 1)).cuda().double()), axis=1).T).T
    p3d = p3d[:,:3]
    return p3d
//  [python]
def forward(self, X, Y):
        """Override abstract function."""

        n, k = X.shape[:2]
        device = X.device

        Z = torch.zeros(n, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device)
        K = K_matrix(X, X)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        Q = torch.solve(Z, L)[0]
        #Q = torch.linalg.solve(L, Z)
        # return W and A.
        return Q[:, :k], Q[:, k:]
//  [python]
def forward(self, X, Y):
        """Override abstract function."""

        n, k = X.shape[:2]
        device = X.device

        Z = torch.zeros(n, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device)
        K = K_matrix(X, X)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        Q = torch.linalg.solve(L, Z)
        # return W and A.
        return Q[:, :k], Q[:, k:]
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        batch_size = inputs.shape[0]
        outputs = inputs - self.bias
        # outputs, lu = torch.solve(outputs.t(), self._weight)
        outputs, lu = torch.solve(outputs.t(), self._weight.weight)
        outputs = outputs.t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * torch.ones(batch_size, device=inputs.device)
        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        batch_size = inputs.shape[0]
        outputs = inputs - self.bias
        outputs, lu = torch.linalg.solve(self._weight.weight, outputs.t())
        outputs = outputs.t()
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * torch.ones(batch_size, device=inputs.device)
        return outputs, logabsdet
//  [python]
def update(self, mean, covariance, measurement):
        """Run Kalman filter correction step.

        Parameters
        ----------
        mean : ndarray
            The predicted state's mean vector (*, 8).
        covariance : ndarray
            The state's covariance matrix (*,8,8).
        measurement : ndarray
            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
            is the center position, a the aspect ratio, and h the height of the
            bounding box.

        Returns
        -------
        (ndarray, ndarray)
            Returns the measurement-corrected state distribution.

        """

        projected_mean, projected_cov = self.project(mean, covariance)

        # Unfortunately, cholesky will randomly throw CUDA ERROR under linux GPU environment,
        # this error is from pytorch1.2 to 1.5
        # chol_factor = torch.cholesky(projected_cov, upper=False)
        # kalman_gain = torch.cholesky_solve(torch.matmul(covariance, self._update_mat).permute(0, 2, 1),
        #                                    chol_factor,
        #                                    upper=False).permute(0, 2, 1)

        # The alternative solution is theoretically slower than cholesky_solve
        kalman_gain = torch.solve(torch.matmul(covariance, self._update_mat).permute(0, 2, 1), projected_cov)[
            0].permute(0, 2, 1)

        # (*, 4)
        innovation = measurement.view(-1, 4) - projected_mean

        kalman_gain_t = kalman_gain.permute(0, 2, 1)
        new_mean = mean + torch.bmm(innovation.unsqueeze(1), kalman_gain_t).view(-1, 8)  # (*, 8)
        new_covariance = covariance - torch.matmul(
            torch.matmul(projected_cov.permute(0, 2, 1), kalman_gain_t).permute(0, 2, 1),
            kalman_gain_t)

        return new_mean, new_covariance
//  [python]
def update(self, mean, covariance, measurement):
        """Run Kalman filter correction step.

        Parameters
        ----------
        mean : ndarray
            The predicted state's mean vector (*, 8).
        covariance : ndarray
            The state's covariance matrix (*,8,8).
        measurement : ndarray
            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
            is the center position, a the aspect ratio, and h the height of the
            bounding box.

        Returns
        -------
        (ndarray, ndarray)
            Returns the measurement-corrected state distribution.

        """

        projected_mean, projected_cov = self.project(mean, covariance)

        # Unfortunately, cholesky will randomly throw CUDA ERROR under linux GPU environment,
        # this error is from pytorch1.2 to 1.5
        # chol_factor = torch.cholesky(projected_cov, upper=False)
        # kalman_gain = torch.cholesky_solve(torch.matmul(covariance, self._update_mat).permute(0, 2, 1),
        #                                    chol_factor,
        #                                    upper=False).permute(0, 2, 1)

        # The alternative solution is theoretically slower than cholesky_solve
        kalman_gain = torch.linalg.solve(projected_cov, torch.matmul(covariance, self._update_mat).permute(0, 2, 1)).permute(0, 2, 1)

        # (*, 4)
        innovation = measurement.view(-1, 4) - projected_mean

        kalman_gain_t = kalman_gain.permute(0, 2, 1)
        new_mean = mean + torch.bmm(innovation.unsqueeze(1), kalman_gain_t).view(-1, 8)  # (*, 8)
        new_covariance = covariance - torch.matmul(
            torch.matmul(projected_cov.permute(0, 2, 1), kalman_gain_t).permute(0, 2, 1),
            kalman_gain_t)

        return new_mean, new_covariance
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        batch_size = inputs.shape[0]
        outputs = inputs - self.bias
        # outputs, lu = torch.gesv(outputs.t(), self._weight)  # Linear-system solver.
        outputs, lu = torch.solve(outputs.t(), self._weight)
        outputs = outputs.t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * torch.ones(batch_size, device=inputs.device)
        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        batch_size = inputs.shape[0]
        outputs = inputs - self.bias
        outputs, lu = torch.linalg.solve(self._weight, outputs.t())
        outputs = outputs.t()
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * torch.ones(batch_size, device=inputs.device)
        return outputs, logabsdet
//  [python]
def product_gaussians(mu1, sigma1, mu2, sigma2, beta=0.0):
    # https://math.stackexchange.com/questions/157172/product-of-two-multivariate-gaussians-distributions
    # sigma1 = prior, sigma2 = neural networks
    # \Sigma_1 + \Sigma_2
    K_sum = sigma1 * (1 - beta) + sigma2 * beta

    # \Sigma_3 = \Sigma_1(\Sigma_1 + \Sigma_2)^{-1}\Sigma_2
    K_adapt = sigma1.matmul(K_sum.inv_matmul(sigma2))

    # \mu_3 = \Sigma_3\Sigma_2^{-1}\mu_2
    # \mu_3 = \Sigma_3\Sigma_1^{-1}\mu_1 + \Sigma_3\Sigma_2^{-1}\mu_2
    mean_adapt = K_adapt.mm(torch.solve(mu2.unsqueeze(-1), sigma2)[0]) * (1 - beta)
    return mean_adapt, K_adapt
//  [python]
def product_gaussians(mu1, sigma1, mu2, sigma2, beta=0.0):
    # https://math.stackexchange.com/questions/157172/product-of-two-multivariate-gaussians-distributions
    # sigma1 = prior, sigma2 = neural networks
    # \Sigma_1 + \Sigma_2
    K_sum = sigma1 * (1 - beta) + sigma2 * beta

    # \Sigma_3 = \Sigma_1(\Sigma_1 + \Sigma_2)^{-1}\Sigma_2
    K_adapt = sigma1.matmul(K_sum.inv_matmul(sigma2))

    # \mu_3 = \Sigma_3\Sigma_2^{-1}\mu_2
    # \mu_3 = \Sigma_3\Sigma_1^{-1}\mu_1 + \Sigma_3\Sigma_2^{-1}\mu_2
    mean_adapt = K_adapt.mm(torch.linalg.solve(sigma2, mu2.unsqueeze(-1))) * (1 - beta)
    return mean_adapt, K_adapt
//  [python]
def unproject2(p2d, Kdist, Mdist, distCoeffs, device):
    # st()
    # pn2d = torch.solve(Kdist, torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0]), device=device)), axis=0)).T
    
    pn2d, _ = torch.solve(torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0]), device=device).double()), axis=0), Kdist)
    pn2d = pn2d.T

    # pn2d = torch.linalg.solve(Kdist, torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0]), device=device)), axis=0)).T

    # print(f'pn2d.mean() {pn2d.mean()}')

    # k = torch.cat([distCoeffs[:5], np.zeros(12-5)], axis=0)
    k = torch.cat((distCoeffs[:len(distCoeffs)], torch.zeros(12-len(distCoeffs), device=device).double()), axis=0)
    x0 = pn2d[:,0]
    y0 = pn2d[:,1]
    x = x0
    y = y0

    # print(f'x {x}')
    # print(f'y {y}')

    for iter in range(5):
        r2 = x*x + y*y
        # st()
        icdist = (1 + ((k[7]*r2 + k[6])*r2 + k[5])*r2)/(1 + ((k[4]*r2 + k[1])*r2 + k[0])*r2)
        deltaX = 2*k[2]*x*y + k[3]*(r2 + 2*x*x)+ k[8]*r2+k[9]*r2*r2
        deltaY = k[2]*(r2 + 2*y*y) + 2*k[3]*x*y+ k[10]*r2+k[11]*r2*r2
        x = (x0 - deltaX)*icdist
        y = (y0 - deltaY)*icdist
    
    pn2d = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), axis=1)
    
    if vis_flag:
        print(f'pn2d.mean() {pn2d.mean()}')
        print(f'p2d.mean() {p2d.mean()}')
        print(f'(p2d[:,2:3]*0.001).mean() {(p2d[:,2:3]*0.001).mean()}')
        print(f'(pn2d * (p2d[:,2:3]*0.001)).mean() {(pn2d * (p2d[:,2:3]*0.001)).mean()}')
        print(f'(p2d[:,2:3]*0.001).mean() {(p2d[:,2:3]*0.001).mean()}')
        print(f'pn2d {pn2d}')
    
    p3d = torch.cat([pn2d * (p2d[:,2:3]*0.001), p2d[:,2:3]*0.001], axis=1)
    # p3d = np.concatenate([pn2d * (p2d[:,2:3]*0.001), p2d[:,2:3]*0.001], axis=1)
    
    if vis_flag:
        print(f'p3d.mean() {p3d.mean()}')

    # st()
    
    # print(p3d.mean(axis=0))
    # st()
    # p3d = np.matmul(np.linalg.inv(np.row_stack([Mdist, [0,0,0,1]])), torch.cat([p3d, np.ones((p3d.shape[0], 1))], axis=1).T).T
    # print(p3d.mean(axis=0))
    p3d = p3d[:,:3]
    # p3d = (inv([Mdist;[0 0 0 1]])*[p3d ones(size(p3d,1),1)]')'
    # p3d = (inv([cam.Mdist;[0 0 0 1]])*[p3d ones(size(p3d,1),1)]')';
    # p3d = [bsxfun(@times, pn2d, p2d(:,3)*0.001) p2d(:,3)*0.001];
    ### 估计没用
    # if pn2d.shape[1] == 1:
    #     pn2d = pn2d.T
    # st()
    # pn2d = [x y]
    # for iter=1:5
    #     r2 = x.*x + y.*y;
    #     icdist = (1 + ((k(1+7)*r2 + k(1+6)).*r2 + k(1+5)).*r2)./(1 + ((k(1+4)*r2 + k(1+1)).*r2 + k(1+0)).*r2);
    #     deltaX = 2*k(1+2)*x.*y + k(1+3)*(r2 + 2*x.*x)+ k(1+8)*r2+k(1+9)*r2.*r2;
    #     deltaY = k(1+2)*(r2 + 2*y.*y) + 2*k(1+3)*x.*y+ k(1+10)*r2+k(1+11)*r2.*r2;
    #     x = (x0 - deltaX).*icdist;
    #     y = (y0 - deltaY).*icdist;
    # end
    # k = [reshape(cam.distCoeffs(1:5),[],1); zeros(12-5,1)];
    # pn2d = ( cam.Kdist \ [p2d(:,1:2)'; ones(1, size(p2d,1))] )';
    return p3d
//  [python]
def unproject2(p2d, Kdist, Mdist, distCoeffs, device):
    pn2d, _ = torch.linalg.solve(Kdist, torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0]), device=device).double()), axis=0))
    pn2d = pn2d.T

    k = torch.cat((distCoeffs[:len(distCoeffs)], torch.zeros(12-len(distCoeffs), device=device).double()), axis=0)
    x0 = pn2d[:,0]
    y0 = pn2d[:,1]
    x = x0
    y = y0

    for iter in range(5):
        r2 = x*x + y*y
        icdist = (1 + ((k[7]*r2 + k[6])*r2 + k[5])*r2)/(1 + ((k[4]*r2 + k[1])*r2 + k[0])*r2)
        deltaX = 2*k[2]*x*y + k[3]*(r2 + 2*x*x)+ k[8]*r2+k[9]*r2*r2
        deltaY = k[2]*(r2 + 2*y*y) + 2*k[3]*x*y+ k[10]*r2+k[11]*r2*r2
        x = (x0 - deltaX)*icdist
        y = (y0 - deltaY)*icdist
    
    pn2d = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), axis=1)
    
    if vis_flag:
        print(f'pn2d.mean() {pn2d.mean()}')
        print(f'p2d.mean() {p2d.mean()}')
        print(f'(p2d[:,2:3]*0.001).mean() {(p2d[:,2:3]*0.001).mean()}')
        print(f'(pn2d * (p2d[:,2:3]*0.001)).mean() {(pn2d * (p2d[:,2:3]*0.001)).mean()}')
        print(f'(p2d[:,2:3]*0.001).mean() {(p2d[:,2:3]*0.001).mean()}')
        print(f'pn2d {pn2d}')
    
    p3d = torch.cat([pn2d * (p2d[:,2:3]*0.001), p2d[:,2:3]*0.001], axis=1)
    
    if vis_flag:
        print(f'p3d.mean() {p3d.mean()}')

    p3d = p3d[:,:3]
    return p3d
//  [python]
def inverse(self, z):
        '''x = D^{-1}z - D^{-1}UK^{-1}VD^{-1}z'''
        if self.bias is not None: z = z - self.bias
        VDiz = torch.einsum('rd,bd->br', self.V, z / self.d)
        KiVDiz = torch.solve(VDiz.t(), self.K)[0].t()
        UKiVDiz = torch.einsum('dr,br->bd', self.U, KiVDiz)
        x = (z - UKiVDiz) / self.d
        return x
//  [python]
def inverse(self, z):
        '''x = D^{-1}z - D^{-1}UK^{-1}VD^{-1}z'''
        if self.bias is not None: z = z - self.bias
        VDiz = torch.einsum('rd,bd->br', self.V, z / self.d)
        KiVDiz = torch.linalg.solve(self.K, VDiz.t()).t()
        UKiVDiz = torch.einsum('dr,br->bd', self.U, KiVDiz)
        x = (z - UKiVDiz) / self.d
        return x
//  [python]
def _build_predict(self, Xnew, full_cov=False):
        """
        Compute the mean and variance of the latent function at some new points
        Xnew.
        """
        _, _, Luu, L, _, _, gamma = self._common_terms()
        Kus = self.kern.K(self.Z.get(), Xnew)  # size  M x Xnew

        w, _ = torch.solve(Kus, Luu)  # size M x Xnew

        tmp, _ = torch.solve(gamma, L.t())
        mean = torch.matmul(w.t(), tmp) + self.mean_function(Xnew)
        intermediateA, _ = torch.solve(w, L)

        if full_cov:
            var = self.kern.K(Xnew) - torch.matmul(w.t(), w) + torch.matmul(intermediateA.t(), intermediateA)
            var = torch.unsqueeze(2).expand(-1, -1, self.Y.size(1))
        else:
            var = self.kern.Kdiag(Xnew) - (w**2).sum(0) + (intermediateA**2).sum(0)  # size Xnew,
            var = torch.unsuqeeze(2).expand(-1, self.Y.size(1))

        return mean, var
//  [python]
def _build_predict(self, Xnew, full_cov=False):
        """
        Compute the mean and variance of the latent function at some new points
        Xnew.
        """
        _, _, Luu, L, _, _, gamma = self._common_terms()
        Kus = self.kern.K(self.Z.get(), Xnew)  # size  M x Xnew

        w, _ = torch.linalg.solve(Luu, Kus)  # size M x Xnew

        tmp, _ = torch.linalg.solve(L.t(), gamma)
        mean = torch.matmul(w.t(), tmp) + self.mean_function(Xnew)
        intermediateA, _ = torch.linalg.solve(L, w)

        if full_cov:
            var = self.kern.K(Xnew) - torch.matmul(w.t(), w) + torch.matmul(intermediateA.t(), intermediateA)
            var = torch.unsqueeze(var, 2).expand(-1, -1, self.Y.size(1))
        else:
            var = self.kern.Kdiag(Xnew) - (w**2).sum(0) + (intermediateA**2).sum(0)  # size Xnew,
            var = torch.unsqueeze(var, 2).expand(-1, -1, self.Y.size(1))

        return mean, var
//  [python]
def pspmat(self, src, dst):
        # B, 4, 2
        B, _, _ = src.size()
        s = torch.cat([
            torch.cat([src, torch.ones((B, 4, 1), device='cuda'), torch.zeros((B, 4, 3), device='cuda'), 
                        -dst[..., 0 : 1] * src[..., 0 : 1], -dst[..., 0 : 1] * src[..., 1 : 2]], dim=2),
            torch.cat([torch.zeros((B, 4, 3), device='cuda'), src, torch.ones((B, 4, 1), device='cuda'), 
                        -dst[..., 1 : 2] * src[..., 0 : 1], -dst[..., 1 : 2] * src[..., 1 : 2]], dim=2)
        ], dim=1)
        t = torch.cat([dst[..., 0 : 1], dst[..., 1 : 2]], dim=1)
        # M = s.inverse() @ t
        M = torch.solve(t, s)[0]
        # M is B 8 1
        return M
//  [python]
def pspmat(self, src, dst):
        # B, 4, 2
        B, _, _ = src.size()
        s = torch.cat([
            torch.cat([src, torch.ones((B, 4, 1), device='cuda'), torch.zeros((B, 4, 3), device='cuda'), 
                        -dst[..., 0 : 1] * src[..., 0 : 1], -dst[..., 0 : 1] * src[..., 1 : 2]], dim=2),
            torch.cat([torch.zeros((B, 4, 3), device='cuda'), src, torch.ones((B, 4, 1), device='cuda'), 
                        -dst[..., 1 : 2] * src[..., 0 : 1], -dst[..., 1 : 2] * src[..., 1 : 2]], dim=2)
        ], dim=1)
        t = torch.cat([dst[..., 0 : 1], dst[..., 1 : 2]], dim=1)
        # M = s.inverse() @ t
        M = torch.linalg.solve(s, t)
        # M is B 8 1
        return M
//  [python]
def gesv(b, A, out=None):
    r"""This function returns the solution to the system of linear equations represented
    by :math:`AX = B` and the LU factorization of A, in order as a tuple `X, LU`.

    For more information regarding :func:`torch.gesv`, please check :func:`torch.solve`.

    .. warning::
        :func:`torch.gesv` is deprecated in favour of :func:`torch.solve` and will be removed in the
        next release. Please use :func:`torch.solve` instead.
    """
    warnings.warn("torch.gesv is deprecated in favour of torch.solve and will be removed in the "
                  "next release. Please use torch.solve instead.", stacklevel=2)
    return torch.solve(b, A, out=out)
//  [python]
def gesv(b, A, out=None):
    r"""This function returns the solution to the system of linear equations represented
    by :math:`AX = B` and the LU factorization of A, in order as a tuple `X, LU`.

    For more information regarding :func:`torch.gesv`, please check :func:`torch.solve`.

    .. warning::
        :func:`torch.gesv` is deprecated in favour of :func:`torch.solve` and will be removed in the
        next release. Please use :func:`torch.solve` instead.
    """
    warnings.warn("torch.gesv is deprecated in favour of torch.solve and will be removed in the "
                  "next release. Please use torch.solve instead.", stacklevel=2)
    return torch.linalg.solve(A, b, out=out)
//  [python]
def _compute_alpha(points, bases):
        # Compute weights (..., N, 4) corresponded to bases (..., N, 4)
        # for points (..., N, 3). Check equation 1 in paper for details.
        points, bases = cart2homo(points), cart2homo(bases)
        return torch.solve(B=points, A=bases)[0]
//  [python]
def _compute_alpha(points, bases):
        # Compute weights (..., N, 4) corresponded to bases (..., N, 4)
        # for points (..., N, 3). Check equation 1 in paper for details.
        points, bases = cart2homo(points), cart2homo(bases)
        return torch.linalg.solve(A=bases, B=points, left=False)
//  [python]
def fit_full_poly(self, epsilon=1, k=2, lmbd=0, target='hypo', fkine=None):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        self.poly_kernel = kernel.Polyharmonic(k=k, epsilon=epsilon)
        phi = self.poly_kernel(X, X)
        phi.fill_diagonal_(lmbd)
        print(phi.shape)
        l1 = torch.cat([phi, X, torch.ones((len(X), 1))], dim=1)
        l2 = torch.cat([X.T, torch.zeros((X.shape[1], X.shape[1]+1))], dim=1)
        l3 = torch.cat([torch.ones((1, len(X))), torch.zeros(1, X.shape[1]+1)], dim=1)
        # print([l.shape for l in [l1, l2, l3]])
        L = torch.cat([l1, l2, l3], dim=0)
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.poly_nodes = torch.solve(
            torch.cat([y, torch.zeros(X.shape[1]+1)], dim=0).reshape(-1, 1),
            L).solution.reshape(-1)
//  [python]
def fit_full_poly(self, epsilon=1, k=2, lmbd=0, target='hypo', fkine=None):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        self.poly_kernel = kernel.Polyharmonic(k=k, epsilon=epsilon)
        phi = self.poly_kernel(X, X)
        phi.fill_diagonal_(lmbd)
        print(phi.shape)
        l1 = torch.cat([phi, X, torch.ones((len(X), 1))], dim=1)
        l2 = torch.cat([X.T, torch.zeros((X.shape[1], X.shape[1]+1))], dim=1)
        l3 = torch.cat([torch.ones((1, len(X))), torch.zeros(1, X.shape[1]+1)], dim=1)
        # print([l.shape for l in [l1, l2, l3]])
        L = torch.cat([l1, l2, l3], dim=0)
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.poly_nodes = torch.linalg.solve(
            torch.cat([y, torch.zeros(X.shape[1]+1)], dim=0).reshape(-1, 1),
            L).reshape(-1)
//  [python]
def twoStageOptimizer(u, weight, A, b, iteration, J, device):
    """
        <weight, A> @ u.T = <weight, b>
    """

    # get cumulative <weight, A> and <weight, b>
    A_ = torch.zeros(A[0].shape).to(device)
    b_ = torch.zeros(b[0].shape).to(device)
    for i in range(len(A)):
        A_ += weight[i] * A[i]
        b_ += weight[i] * b[i]

    # two stage optimization
    if (iteration < 1) and (J == 1):
        u = jacobi(u, A_, b_, device)
    else:
        u = torch.solve(b_, A_ + torch.eye(A_.shape[1]).to(device) * 1e-5)[0]
        # L = torch.cholesky(A_ + torch.eye(A_.shape[1]).to(device) * 1e-5)
        # y = torch.triangular_solve(b_, L.T, upper=True)[0]
        # u1 = torch.triangular_solve(y, L, upper=False)[0]
    return u
//  [python]
def twoStageOptimizer(u, weight, A, b, iteration, J, device):
    """
        <weight, A> @ u.T = <weight, b>
    """

    # get cumulative <weight, A> and <weight, b>
    A_ = torch.zeros(A[0].shape).to(device)
    b_ = torch.zeros(b[0].shape).to(device)
    for i in range(len(A)):
        A_ += weight[i] * A[i]
        b_ += weight[i] * b[i]

    # two stage optimization
    if (iteration < 1) and (J == 1):
        u = jacobi(u, A_, b_, device)
    else:
        u = torch.linalg.solve(A_ + torch.eye(A_.shape[1]).to(device) * 1e-5, b_)
        # L = torch.cholesky(A_ + torch.eye(A_.shape[1]).to(device) * 1e-5)
        # y = torch.triangular_solve(b_, L.T, upper=True)[0]
        # u1 = torch.triangular_solve(y, L, upper=False)[0]
    return u
//  [python]
def forward(self, src, dst):
        # src and dst are B.n.2
        B, n, _ = src.size()
        # B.n.1.2
        delta = src.unsqueeze(2)
        delta = delta - delta.permute(0, 2, 1, 3)
        # B.n.n
        K = delta.norm(dim=3)
        # Rsq = torch.sum(delta**2, dim=3)
        # Rsq += torch.eye(n, device='cuda')
        # Rsq[Rsq == 0] = 1.
        # K = 0.5 * Rsq * torch.log(Rsq)
        # c = -150
        # K = torch.exp(c * Rsq)
        # K = torch.abs(Rsq - 0.5) - 0.5
        # WARNING: TORCH.SQRT HAS NAN GRAD AT 0
        # K = torch.sqrt(Rsq)
        # print(K)
        # K[torch.isnan(K)] = 0.
        P = torch.cat((torch.ones((B, n, 1), device='cuda'), src), 2)
        L = torch.cat((K, P), 2)
        t = torch.cat((P.permute(0, 2, 1), torch.zeros((B, 3, 3), device='cuda')), 2)
        L = torch.cat((L, t), 1)
        # LInv = L.inverse()
        # # wv is B.n+3.2
        # wv = torch.matmul(LInv, torch.cat((dst, torch.zeros((B, 3, 2), device='cuda')), 1))
        # the above implementation has stability problem near the boundaries
        wv = torch.solve(torch.cat((dst, torch.zeros((B, 3, 2), device='cuda')), 1), L)[0]

        # get the grid sampler
        s = self.gs.size(1)
        gs = self.gs
        delta = gs.unsqueeze(2)   
        delta = delta - src.unsqueeze(1)
        K = delta.norm(dim=3)
        # Rsq = torch.sum(delta**2, dim=3)
        # K = torch.exp(c * Rsq)
        # Rsq[Rsq == 0] = 1.
        # K =  0.5 * Rsq * torch.log(Rsq)
        # K = torch.abs(Rsq - 0.5) - 0.5
        # K = torch.sqrt(Rsq)
        # K[torch.isnan(K)] = 0.
        gs = gs.expand(B, -1, -1)
        P = torch.cat((torch.ones((B, s, 1), device='cuda'), gs), 2)
        L = torch.cat((K, P), 2)
        gs = torch.matmul(L, wv)
        return gs.reshape(B, self.sz, self.sz, 2).permute(0, 3, 1, 2)
//  [python]
def forward(self, src, dst):
        # src and dst are B.n.2
        B, n, _ = src.size()
        # B.n.1.2
        delta = src.unsqueeze(2)
        delta = delta - delta.permute(0, 2, 1, 3)
        # B.n.n
        K = delta.norm(dim=3)
        # Rsq = torch.sum(delta**2, dim=3)
        # Rsq += torch.eye(n, device='cuda')
        # Rsq[Rsq == 0] = 1.
        # K = 0.5 * Rsq * torch.log(Rsq)
        # c = -150
        # K = torch.exp(c * Rsq)
        # K = torch.abs(Rsq - 0.5) - 0.5
        # WARNING: TORCH.SQRT HAS NAN GRAD AT 0
        # K = torch.sqrt(Rsq)
        # print(K)
        # K[torch.isnan(K)] = 0.
        P = torch.cat((torch.ones((B, n, 1), device='cuda'), src), 2)
        L = torch.cat((K, P), 2)
        t = torch.cat((P.permute(0, 2, 1), torch.zeros((B, 3, 3), device='cuda')), 2)
        L = torch.cat((L, t), 1)
        # LInv = L.inverse()
        # # wv is B.n+3.2
        # wv = torch.matmul(LInv, torch.cat((dst, torch.zeros((B, 3, 2), device='cuda')), 1))
        # the above implementation has stability problem near the boundaries
        wv = torch.linalg.solve(L, torch.cat((dst, torch.zeros((B, 3, 2), device='cuda')), 1))

        # get the grid sampler
        s = self.gs.size(1)
        gs = self.gs
        delta = gs.unsqueeze(2)   
        delta = delta - src.unsqueeze(1)
        K = delta.norm(dim=3)
        # Rsq = torch.sum(delta**2, dim=3)
        # K = torch.exp(c * Rsq)
        # Rsq[Rsq == 0] = 1.
        # K =  0.5 * Rsq * torch.log(Rsq)
        # K = torch.abs(Rsq - 0.5) - 0.5
        # K = torch.sqrt(Rsq)
        # K[torch.isnan(K)] = 0.
        gs = gs.expand(B, -1, -1)
        P = torch.cat((torch.ones((B, s, 1), device='cuda'), gs), 2)
        L = torch.cat((K, P), 2)
        gs = torch.matmul(L, wv)
        return gs.reshape(B, self.sz, self.sz, 2).permute(0, 3, 1, 2)
//  [python]
def _muygps_posterior_mean(
    K: torch.ndarray,
    Kcross: torch.ndarray,
    batch_nn_targets: torch.ndarray,
    **kwargs,
) -> torch.ndarray:
    batch_count, nn_count, response_count = batch_nn_targets.shape
    responses = Kcross.reshape(batch_count, 1, nn_count) @ torch.solve(
        batch_nn_targets, K
    )[0]
    return responses.reshape(batch_count, response_count)
//  [python]
def _muygps_posterior_mean(
    K: torch.ndarray,
    Kcross: torch.ndarray,
    batch_nn_targets: torch.ndarray,
    **kwargs,
) -> torch.ndarray:
    batch_count, nn_count, response_count = batch_nn_targets.shape
    responses = Kcross.reshape(batch_count, 1, nn_count) @ torch.linalg.solve(
        K, batch_nn_targets
    )
    return responses.reshape(batch_count, response_count)
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a, return_LU=False)
        else:
            return matmul(inverse(a), b)
    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.solve(b, a)[0]
        else:
            return matmul(inverse(a), b)
    else:
        if is_torch_1_8_plus:
            return torch.solve(b, a)[0]
        else:
            return torch.solve(b, a)[0]
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a, return_LU=False)
        else:
            return matmul(inverse(a), b)
    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.linalg.solve(a, b)
        else:
            return matmul(inverse(a), b)
    else:
        if is_torch_1_8_plus:
            return torch.linalg.solve(a, b)
        else:
            return torch.solve(b, a)[0]
//  [python]
def find_intersection_2D(line1, line2, eps : float = 1e-7, verify : bool = False):
    '''
    lines1: batch_size x 3
    lines2: batch_size x 3
    '''
    
    Ab = torch.cat(
        (torch.reshape(line1,shape=(line1.shape[0],1,line1.shape[1])),
         torch.reshape(line2,shape=(line2.shape[0],1,line2.shape[1]))),
        dim=1
    )
    
    # Ax + b = 0 -> Ax = -b
    v = torch.solve(Ab[...,:-1], -Ab[...,-1])[0]
    
    flag = False
    if verify:
        
        flag = torch.allclose(
            torch.bmm(Ab[...,:-1],v[...,None]),
            -Ab[...,-1][...,None],
            atol=eps,
            rtol=0.,
            equal_nan=False
        )
    
    return v,flag
//  [python]
def find_intersection_2D(line1, line2, eps : float = 1e-7, verify : bool = False):
    '''
    lines1: batch_size x 3
    lines2: batch_size x 3
    '''
    
    Ab = torch.cat(
        (torch.reshape(line1,shape=(line1.shape[0],1,line1.shape[1])),
         torch.reshape(line2,shape=(line2.shape[0],1,line2.shape[1]))),
        dim=1
    )
    
    # Ax + b = 0 -> Ax = -b
    v = torch.linalg.solve(Ab[...,:-1], -Ab[...,-1])
    
    flag = False
    if verify:
        
        flag = torch.allclose(
            torch.bmm(Ab[...,:-1],v[...,None]),
            -Ab[...,-1][...,None],
            atol=eps,
            rtol=0.,
            equal_nan=False
        )
    
    return v,flag
//  [python]
def even_intergrid(n: int):
    """Calculate the coefficients of taylar expansion of even intergrid

    Args:
        n (int): The spatial order of the laplace operator.
    """
    def fact(n):
        if n == 0:
            return 1
        else:
            return n * fact(n-1)
    
    MAX = n // 2
    matrix = torch.zeros((MAX, MAX))
    constant_one = torch.zeros((MAX,))
    for i in range(MAX):
        for j in range(MAX):
            matrix[i, j] = (j + 1)**(2*(i+1) - 2)
        if i != ((2+1)//2 - 1):
            constant_one[i] = 0
        else:
            constant_one[i] = (1**(2*(i+1)-2) * fact(2*(i+1)-2))

    C = torch.solve(constant_one.unsqueeze(1), matrix)[0]
    constant = C.squeeze() / torch.tensor([(i+1)**2 for i in range(MAX)], dtype=torch.float32)
    return constant
//  [python]
def even_intergrid(n: int):
    """Calculate the coefficients of taylar expansion of even intergrid

    Args:
        n (int): The spatial order of the laplace operator.
    """
    def fact(n):
        if n == 0:
            return 1
        else:
            return n * fact(n-1)
    
    MAX = n // 2
    matrix = torch.zeros((MAX, MAX))
    constant_one = torch.zeros((MAX,))
    for i in range(MAX):
        for j in range(MAX):
            matrix[i, j] = (j + 1)**(2*(i+1) - 2)
        if i != ((2+1)//2 - 1):
            constant_one[i] = 0
        else:
            constant_one[i] = (1**(2*(i+1)-2) * fact(2*(i+1)-2))

    C = torch.linalg.solve(matrix, constant_one.unsqueeze(1))
    constant = C.squeeze() / torch.tensor([(i+1)**2 for i in range(MAX)], dtype=torch.float32)
    return constant
//  [python]
def _linear_solve(self, jac, residual):
        return torch.solve(residual, jac)[1]
//  [python]
def _linear_solve(self, jac, residual):
        return torch.linalg.solve(jac, residual)
//  [python]
def state_and_cov_update(self, Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R):
        S = H.mm(P).mm(H.t()) + R
        K = (torch.solve(H.t().mm(P).t(), S)[0]).t()
        dx = K.mv(r)

        dR, dxi = self.sen3exp(dx[:9])
        dv = dxi[:, 0]
        dp = dxi[:, 1]
        Rot_up = dR.mm(Rot)
        v_up = dR.mv(v) + dv
        p_up = dR.mv(p) + dp

        b_omega_up = b_omega + dx[9:12]
        b_acc_up = b_acc + dx[12:15]

        dR = self.so3exp(dx[15:18])
        Rot_c_i_up = dR.mm(Rot_c_i)
        t_c_i_up = t_c_i + dx[18:21]

        I_KH = self.IdP - K.mm(H)
        P_up = I_KH.mm(P).mm(I_KH.t()) + K.mm(R).mm(K.t())
        P_up = (P_up + P_up.t())/2
        return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
//  [python]
def state_and_cov_update(self, Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R):
        S = H.mm(P).mm(H.t()) + R
        K = (torch.linalg.solve(S, (P.mm(H.t())).t())).t()
        dx = K.mv(r)

        dR, dxi = self.sen3exp(dx[:9])
        dv = dxi[:, 0]
        dp = dxi[:, 1]
        Rot_up = dR.mm(Rot)
        v_up = dR.mv(v) + dv
        p_up = dR.mv(p) + dp

        b_omega_up = b_omega + dx[9:12]
        b_acc_up = b_acc + dx[12:15]

        dR = self.so3exp(dx[15:18])
        Rot_c_i_up = dR.mm(Rot_c_i)
        t_c_i_up = t_c_i + dx[18:21]

        I_KH = self.IdP - K.mm(H)
        P_up = I_KH.mm(P).mm(I_KH.t()) + K.mm(R).mm(K.t())
        P_up = (P_up + P_up.t())/2
        return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
//  [python]
def __init__(self, basis: torch.Tensor, tol: float = 1e-6) -> None:
        """Initialize Wigner-Seitz lattice.

        Parameters
        ----------
        basis
            Basis vectors (i.e. lattice vectors in real-space or reciprocal space)
        """

        self.basis = basis
        self.tol = tol
        self.iFaces, self.faces = get_grid(basis)
        faces = self.faces
        iFaces = self.iFaces

        # Down-select to planes that lie on WS boundary:
        touchWS = torch.where(self.ws_boundary_distance(0.5 * faces) < tol)[0]
        faces = faces[touchWS]
        iFaces = iFaces[touchWS]
        # Down-select to WS faces with finite area:
        finiteArea = (self.ws_plane_count(0.5 * faces) == 1)
        faces = faces[finiteArea]  # must be only face with 0 distance
        iFaces = iFaces[finiteArea]
        n_faces = len(faces)
        assert 6 <= n_faces <= 14
        assert n_faces % 2 == 0

        # Compute vertices by 3-way intersection of planes
        i_face_triplets = get_combinations(n_faces, 3)
        Lhs = 0.5 * faces[i_face_triplets]  # n_triplets x 3 x 3
        intersect = (torch.abs(torch.linalg.det(Lhs)) > tol)
        Lhs = Lhs[intersect]
        rhs = (Lhs ** 2).sum(dim=-1)
        vertices = torch.solve(rhs, Lhs)[0]

        # Down-select vertices to WS boundary:
        on_boundary = (torch.abs(self.ws_boundary_distance(vertices)) < tol)
        vertices = vertices[on_boundary]
        vertices = weld_points(vertices, tol)
        self.vertices = vertices
        n_vertices = len(vertices)

        # Find edges of WS boundary:
        i_vertex_pairs = get_combinations(n_vertices, 2)
        edge_midpoints = vertices[i_vertex_pairs].mean(axis=1)
        on_boundary = (self.ws_plane_count(edge_midpoints) == 2)
        edges = i_vertex_pairs[on_boundary]
        self.edges = edges
//  [python]
def __init__(self, basis: torch.Tensor, tol: float = 1e-6) -> None:
        """Initialize Wigner-Seitz lattice.

        Parameters
        ----------
        basis
            Basis vectors (i.e. lattice vectors in real-space or reciprocal space)
        """

        self.basis = basis
        self.tol = tol
        self.iFaces, self.faces = get_grid(basis)
        faces = self.faces
        iFaces = self.iFaces

        # Down-select to planes that lie on WS boundary:
        touchWS = torch.where(self.ws_boundary_distance(0.5 * faces) < tol)[0]
        faces = faces[touchWS]
        iFaces = iFaces[touchWS]
        # Down-select to WS faces with finite area:
        finiteArea = (self.ws_plane_count(0.5 * faces) == 1)
        faces = faces[finiteArea]  # must be only face with 0 distance
        iFaces = iFaces[finiteArea]
        n_faces = len(faces)
        assert 6 <= n_faces <= 14
        assert n_faces % 2 == 0

        # Compute vertices by 3-way intersection of planes
        i_face_triplets = get_combinations(n_faces, 3)
        Lhs = 0.5 * faces[i_face_triplets]  # n_triplets x 3 x 3
        intersect = (torch.abs(torch.linalg.det(Lhs)) > tol)
        Lhs = Lhs[intersect]
        rhs = (Lhs ** 2).sum(dim=-1)
        vertices = torch.linalg.solve(Lhs, rhs)

        # Down-select vertices to WS boundary:
        on_boundary = (torch.abs(self.ws_boundary_distance(vertices)) < tol)
        vertices = vertices[on_boundary]
        vertices = weld_points(vertices, tol)
        self.vertices = vertices
        n_vertices = len(vertices)

        # Find edges of WS boundary:
        i_vertex_pairs = get_combinations(n_vertices, 2)
        edge_midpoints = vertices[i_vertex_pairs].mean(axis=1)
        on_boundary = (self.ws_plane_count(edge_midpoints) == 2)
        edges = i_vertex_pairs[on_boundary]
        self.edges = edges
//  [python]
def _muygps_diagonal_variance(
    K: torch.ndarray,
    Kcross: torch.ndarray,
    **kwargs,
) -> torch.ndarray:
    batch_count, nn_count = Kcross.shape
    return 1 - torch.sum(
        Kcross
        * torch.solve(
            Kcross.reshape(batch_count, nn_count, 1),
            K,
        ).solution.reshape(batch_count, nn_count),
        axis=1,
    )
//  [python]
def _muygps_diagonal_variance(
    K: torch.ndarray,
    Kcross: torch.ndarray,
    **kwargs,
) -> torch.ndarray:
    batch_count, nn_count = Kcross.shape
    return 1 - torch.sum(
        Kcross
        * torch.linalg.solve(
            K, Kcross.reshape(batch_count, nn_count, 1)
        ).reshape(batch_count, nn_count),
        axis=1,
    )
//  [python]
def _muygps_fast_posterior_mean_precompute(
    K: torch.ndarray,
    train_nn_targets_fast: torch.ndarray,
) -> torch.ndarray:
    return torch.solve(train_nn_targets_fast, K)[0]
//  [python]
def _muygps_fast_posterior_mean_precompute(
    K: torch.ndarray,
    train_nn_targets_fast: torch.ndarray,
) -> torch.ndarray:
    return torch.linalg.solve(K, train_nn_targets_fast)
//  [python]
def KL_term(self, mu, L, Kmm):
        Ltril = torch.tril(L)
        hh_expt = torch.matmul(Ltril, Ltril.T) + torch.matmul(mu, mu.T)
        kl = (0.5 * torch.trace(torch.solve(hh_expt, Kmm)[0]) + 0.5 * torch.logdet(Kmm) - 0.5 * torch.sum(torch.log(torch.square(torch.diag(Ltril)))))
        return kl
//  [python]
def KL_term(self, mu, L, Kmm):
        Ltril = torch.tril(L)
        hh_expt = torch.matmul(Ltril, Ltril.T) + torch.matmul(mu, mu.T)
        kl = (0.5 * torch.trace(torch.linalg.solve(Kmm, hh_expt)) + 0.5 * torch.logdet(Kmm) - 0.5 * torch.sum(torch.log(torch.square(torch.diag(Ltril)))))
        return kl
//  [python]
def update_Q_IP(self):
        for m in range(self.n_mic):
            V_FMM = (
                torch.einsum("ftij, ft -> fij", self.XX_FTMM, (1 / self.Y_FTM[..., m]).to(self.TYPE_COMPLEX))
                / self.n_time
            )
            tmp_FM = torch.solve(
                torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None], self.Q_FMM @ V_FMM
            )[0][..., m]
            self.Q_FMM[:, m] = (
                tmp_FM / torch.sqrt(torch.einsum("fi, fij, fj -> f", tmp_FM.conj(), V_FMM, tmp_FM))[:, None]
            ).conj()
//  [python]
def update_Q_IP(self):
        for m in range(self.n_mic):
            V_FMM = (
                torch.einsum("ftij, ft -> fij", self.XX_FTMM, (1 / self.Y_FTM[..., m]).to(self.TYPE_COMPLEX))
                / self.n_time
            )
            tmp_FM = torch.linalg.solve(
                self.Q_FMM @ V_FMM, torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None]
            )[..., m]
            self.Q_FMM[:, m] = (
                tmp_FM / torch.sqrt(torch.einsum("fi, fij, fj -> f", tmp_FM.conj(), V_FMM, tmp_FM))[:, None]
            ).conj()
//  [python]
def solve(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:  # pragma: no cover
    """
    Like torch.linalg.solve, tries to return X
    such that AX=B, with A square.
    """
    if hasattr(torch, "linalg") and hasattr(torch.linalg, "solve"):
        # PyTorch version >= 1.8.0
        return torch.linalg.solve(A, B)

    return torch.solve(B, A).solution
//  [python]
def separate(self, mic_index=MIC_INDEX):
        Y_NFTM = torch.einsum("nft, nfm -> nftm", self.PSD_NFT, self.G_NFM).to(self.TYPE_COMPLEX)
        self.Y_FTM = Y_NFTM.sum(axis=0)
        self.Qx_FTM = torch.einsum("fmi, fti -> ftm", self.Q_FMM, self.X_FTM)
        Qinv_FMM = torch.solve(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None], self.Q_FMM
        )[1]

        self.separated_spec = torch.einsum(
            "fj, ftj, nftj -> nft", Qinv_FMM[:, mic_index], self.Qx_FTM / self.Y_FTM, Y_NFTM
        )
        return self.separated_spec
//  [python]
def separate(self, mic_index=MIC_INDEX):
        Y_NFTM = torch.einsum("nft, nfm -> nftm", self.PSD_NFT, self.G_NFM).to(self.TYPE_COMPLEX)
        self.Y_FTM = Y_NFTM.sum(axis=0)
        self.Qx_FTM = torch.einsum("fmi, fti -> ftm", self.Q_FMM, self.X_FTM)
        Qinv_FMM = torch.linalg.solve(
            self.Q_FMM, torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None]
        )

        self.separated_spec = torch.einsum(
            "fj, ftj, nftj -> nft", Qinv_FMM[:, mic_index], self.Qx_FTM / self.Y_FTM, Y_NFTM
        )
        return self.separated_spec
//  [python]
def compute_rots(self, G_init, G_pose): 
        G_rel = torch.zeros(G_init.shape, dtype=torch.float32, device=self.device) 
        G_rel = torch.solve(G_pose.transpose(2,3), G_init.transpose(2,3))[0].transpose(2,3) 
        R_diff = torch.zeros(G_init.shape, dtype=torch.float32, device=self.device) 
        for i in range(self.jointnum): 
            if i == 0: 
                R_diff[:,i] = G_rel[:,i]
            else: 
                R_diff[:,i] = torch.solve(G_rel[:,i], G_rel[:,self.parents[i]])[0]
        return R_diff
//  [python]
def compute_rots(self, G_init, G_pose): 
        G_rel = torch.zeros(G_init.shape, dtype=torch.float32, device=self.device) 
        G_rel = torch.linalg.solve(G_init.transpose(2,3), G_pose.transpose(2,3)).transpose(2,3) 
        R_diff = torch.zeros(G_init.shape, dtype=torch.float32, device=self.device) 
        for i in range(self.jointnum): 
            if i == 0: 
                R_diff[:,i] = G_rel[:,i]
            else: 
                R_diff[:,i] = torch.linalg.solve(G_rel[:,self.parents[i]], G_rel[:,i])
        return R_diff
//  [python]
def get_sdf_interp(gt_pts_sdfs, pred_pts):
    from pytorch3d.ops import ball_query, knn_points
    gt_pts_sdfs = gt_pts_sdfs.float()
    gt_pts = gt_pts_sdfs[:, :, :3]
    gt_sdfs = gt_pts_sdfs[:, :, 3]
    knn = knn_points(pred_pts, gt_pts, K=8, return_nn=True)
    knn_idx = knn.idx
    knn_pts = knn.knn
    kernel_matrix = linear_plate_kernel(knn_pts, knn_pts)
    closest_sdfs = gt_sdfs[torch.arange(gt_pts_sdfs.shape[0]).unsqueeze(
        1), knn_idx.view(gt_pts_sdfs.shape[0], -1)].view(kernel_matrix.size()[:-1])
    X_mat = torch.solve(closest_sdfs, kernel_matrix)[0]
    closest_sdfs_interpolated = torch.einsum(
        'ijlk,ijk->ijl', linear_plate_kernel(pred_pts.unsqueeze(-2), knn_pts), X_mat)
    return closest_sdfs_interpolated
//  [python]
def get_sdf_interp(gt_pts_sdfs, pred_pts):
    from pytorch3d.ops import ball_query, knn_points
    gt_pts_sdfs = gt_pts_sdfs.float()
    gt_pts = gt_pts_sdfs[:, :, :3]
    gt_sdfs = gt_pts_sdfs[:, :, 3]
    knn = knn_points(pred_pts, gt_pts, K=8, return_nn=True)
    knn_idx = knn.idx
    knn_pts = knn.knn
    kernel_matrix = linear_plate_kernel(knn_pts, knn_pts)
    closest_sdfs = gt_sdfs[torch.arange(gt_pts_sdfs.shape[0]).unsqueeze(
        1), knn_idx.view(gt_pts_sdfs.shape[0], -1)].view(kernel_matrix.size()[:-1])
    X_mat = torch.linalg.solve(kernel_matrix, closest_sdfs)
    closest_sdfs_interpolated = torch.einsum(
        'ijlk,ijk->ijl', linear_plate_kernel(pred_pts.unsqueeze(-2), knn_pts), X_mat)
    return closest_sdfs_interpolated
//  [python]
def predict(self, x):
        """Computes the predicted means and variances at the given location.

        The posterior means and variances are returned if `fit()` has already been
        called otherwise the prior means and variances are returned.

        Args:
            x: The location to predict at.
        Returns
            The predicted means and variances.
        """
        if self.train_x is None:
            return self.mean(x), self.kernel(x, x)

        # Compute means and covariances.
        train_mu, x_mu = self.mean(self.train_x), self.mean(x)
        train_sig = self.kernel(
            self.train_x, self.train_x
        ) + self.noise_var * torch.eye(self.train_x.shape[0])
        x_sig, cross_sig = self.kernel(x, x), self.kernel(self.train_x, x)

        # Compute posterior mean and covariance.
        solved = torch.solve(cross_sig.T, train_sig)[0].T
        mu = x_mu + solved @ (self.train_y - train_mu)
        sig = x_sig - (solved @ cross_sig)

        return mu, sig