jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def estimate_minimal_model(self, pts, weights=None):  # x1 y1 x2 y2
        """Using 5 points to estimate Essential matrix."""
        try:
            pts.shape[1] == self.sample_size
        except Exception as e:
            print(e, "This is not a minimal sample.")

        batch_size, num, _ = pts.shape
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]

        # get the points
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        # Step1: construct the A matrix, A F = 0.
        # 5 equations for 9 variables, A is 5x9 matrix containing epipolar constraints
        # Essential matrix is a linear combination of the 4 vectors spanning the null space of A
        a_59 = torch.ones(x1.shape, device=self.device)
        A_s = torch.stack(
            (torch.mul(x1, x2), torch.mul(x1, y2), x1,
             torch.mul(y1, x2), torch.mul(y1, y2), y1,
             x2, y2, a_59), dim=-1)

        _, _, v = torch.linalg.svd(A_s)#, full_matrices=False) #A_s # eigenvalues in increasing order

        null_space = v[:, -4:, :].transpose(-1, -2)#.float()  # the last four rows

        # use the 4 eigenvectors according to the 4 smallest singular values,
        # E is calculated from 4 basis, E = cx*X + cy*Y + cz*Z + cw*W, up yo common scale = 1
        # X, Y, Z, W = v[:, -1].reshape(3, 3), v[:, -2].reshape(3, 3),
        # v[:, -3].reshape(3, 3), v[:, -4].reshape(3, 3)  # null space
        # null_space_mat= null_space.reshape(3, 3, 4) #X, Y, Z, W
        null_space_mat = null_space.reshape(null_space.shape[0], 3, 3, null_space.shape[-1]).transpose(1, 2)  # X, Y, Z, W
        # Step2: expansion of the constraints:
        # determinant constraint det(E) = 0,
        # trace constraint $2EE^{T}E - trace(EE^{T)}E = 0$
        constraint_mat = self.get_constraint_mat(null_space_mat, batch_size)

        # Step 3: Eliminate part of the matrix to isolate polynomials in z.
        # solve AX=b
        b = constraint_mat[:, :, 10:]
        eliminated_mat = torch.linalg.solve(constraint_mat[:, :, :10], b)

        action_mat = torch.zeros((batch_size, 10, 10), device=self.device)
        action_mat[:, 0:3] = eliminated_mat[:, 0:3]
        action_mat[:, 3] = eliminated_mat[:, 4]
        action_mat[:, 4] = eliminated_mat[:, 5]
        action_mat[:, 5] = eliminated_mat[:, 7]
        action_mat[:, 6, 0] = -torch.ones_like(action_mat[:, 6, 0])
        action_mat[:, 7, 1] = -torch.ones_like(action_mat[:, 7, 1])
        action_mat[:, 8, 3] = -torch.ones_like(action_mat[:, 8, 3])
        action_mat[:, 9, 6] = -torch.ones_like(action_mat[:, 9, 6])

        ee, vv = torch.linalg.eig(action_mat)

        # put the cx, cy, cz back to get a valid essential matrix
        E_models = null_space.matmul(vv.real[:, -4:])#torch.stack(#.real
        E_models = E_models.transpose(-1, -2).reshape(-1, 3, 3).transpose(-1, -2)

        return E_models
//  [python]
def get_heatmap(self, img):
        with torch.no_grad():
            tensor = self.prep(img)[None, ...].to(self.device)
            output = self.model(tensor)
            feature = self.feature['output']
            self._check(feature)
            
            _, _, vT = torch.svd(feature)
            v1 = vT[:, :, 0, :][..., None, :]
            
            cam = feature @ v1.repeat(1, 1, v1.shape[3], 1)
            cam = cam.sum(1)
            cam -= cam.min()
            cam = cam / cam.max() * 255
            cam = cam.cpu().numpy().transpose(1, 2, 0).astype(np.uint8)
            cam = cv2.resize(cam, img.size)
            cam = cv2.applyColorMap(cam, cv2.COLORMAP_JET)        

            if not isinstance(img, np.ndarray):
                img = np.asarray(img)

            overlay = np.uint8(0.6 * img + 0.4 * cam)

        return output, overlay
//  [python]
def get_heatmap(self, img):
        with torch.no_grad():
            tensor = self.prep(img)[None, ...].to(self.device)
            output = self.model(tensor)
            feature = self.feature['output']
            self._check(feature)
            
            _, _, vT = torch.linalg.svd(feature)
            v1 = vT[:, :, 0, :][..., None, :]
            
            cam = feature @ v1.repeat(1, 1, v1.shape[3], 1)
            cam = cam.sum(1)
            cam -= cam.min()
            cam = cam / cam.max() * 255
            cam = cam.cpu().numpy().transpose(1, 2, 0).astype(np.uint8)
            cam = cv2.resize(cam, img.size)
            cam = cv2.applyColorMap(cam, cv2.COLORMAP_JET)        

            if not isinstance(img, np.ndarray):
                img = np.asarray(img)

            overlay = np.uint8(0.6 * img + 0.4 * cam)

        return output, overlay
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(self.out_features, self.in_features, dtype=self.U.dtype, device=self.device)
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones(S.shape[0], dtype=self.U.dtype, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(self.out_features, self.in_features, dtype=self.U.dtype, device=self.device)
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(self.out_features, self.in_features, dtype=self.U.dtype, device=self.device)
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd")  # must use QR decomposition
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones(S.shape[0], dtype=self.U.dtype, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(self.out_features, self.in_features, dtype=self.U.dtype, device=self.device)
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd")  # must use QR decomposition
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def hysime(input, noise, noise_corr):
    """
    HySime: Hyperspectral signal subspace estimation, adapted from [1].

    Parameters
    ----------
    input : torch.Tensor
    noise : torch.Tensor
    noise_corr : torch.Tensor

    Returns
    -------
    signal_subspace_dim : torch.Tensor
    eigs_span_subspace : torch.Tensor
        matrix which columns are the eigenvalues that span the signal subspace

    References
    ----------
    [1] S. Devika and S. M. K. Chaitanya, "Signal estimation of hyperspectral data using HYSIME algorithm," 2016
        International Conference on Research Advances in Integrated Navigation Systems (RAINS), 2016, pp. 1-3,
        doi: 10.1109/RAINS.2016.7764372.
    """
    # [L N] = size(y);
    l, n = input.shape
    # [Ln Nn] = size(n);
    ln, nn = noise.shape
    # [d1 d2] = size(Rn);
    d1, d2 = noise_corr.shape
    if ln != l or nn != n:
        raise RuntimeError("incompatible sizes for noise and input")
    if d1 != d2 or d1 != l:
        print("Bad noise correlation matrix")
        noise_corr = noise @ noise.T / float(n)

    x = input - noise
    # signal correlation matrix estimates
    ry = (input @ input.T) / float(n)
    rx = (x @ x.T) / float(n)
    # eigen values of Rx in decreasing order, equation(15)
    u, s, vh = torch.svd(rx)
    # dx = diag(D);
    # Rn = Rn + sum(diag(Rx)) / L / 10 ^ 5 * eye(L);
    noise_corr += rx.sum() / l / 100000 * torch.eye(l, device=input.device, dtype=input.dtype)
    # Py = diag(E'*Ry*E); %equation (23)
    py = torch.diag(u.T @ ry @ u)
    # Pn = diag(E'*Rn*E); %equation (24)
    pn = torch.diag(u.T @ noise_corr @ u)
    # cost_F = -Py + 2 * Pn; % equation(22)
    cost_f = -py + 2 * pn
    # print(cost_f)
    # kf = sum(cost_F < 0);
    sig_subspace_dim = (cost_f < 0).sum()
    # [dummy, ind_asc] = sort(cost_F, 'ascend');
    _, indices = torch.sort(cost_f)
    # Ek = E(:, ind_asc(1: kf));
    eigs_span_subspace = u[:, indices[:sig_subspace_dim]]
    return sig_subspace_dim, eigs_span_subspace
//  [python]
def hysime(input, noise, noise_corr):
    """
    HySime: Hyperspectral signal subspace estimation, adapted from [1].

    Parameters
    ----------
    input : torch.Tensor
    noise : torch.Tensor
    noise_corr : torch.Tensor

    Returns
    -------
    signal_subspace_dim : torch.Tensor
    eigs_span_subspace : torch.Tensor
        matrix which columns are the eigenvalues that span the signal subspace

    References
    ----------
    [1] S. Devika and S. M. K. Chaitanya, "Signal estimation of hyperspectral data using HYSIME algorithm," 2016
        International Conference on Research Advances in Integrated Navigation Systems (RAINS), 2016, pp. 1-3,
        doi: 10.1109/RAINS.2016.7764372.
    """
    # [L N] = size(y);
    l, n = input.shape
    # [Ln Nn] = size(n);
    ln, nn = noise.shape
    # [d1 d2] = size(Rn);
    d1, d2 = noise_corr.shape
    if ln != l or nn != n:
        raise RuntimeError("incompatible sizes for noise and input")
    if d1 != d2 or d1 != l:
        print("Bad noise correlation matrix")
        noise_corr = noise @ noise.T / float(n)

    x = input - noise
    # signal correlation matrix estimates
    ry = (input @ input.T) / float(n)
    rx = (x @ x.T) / float(n)
    # eigen values of Rx in decreasing order, equation(15)
    u, s, vh = torch.linalg.svd(rx, full_matrices=False)
    # dx = diag(D);
    # Rn = Rn + sum(diag(Rx)) / L / 10 ^ 5 * eye(L);
    noise_corr += rx.sum() / l / 100000 * torch.eye(l, device=input.device, dtype=input.dtype)
    # Py = diag(E'*Ry*E); %equation (23)
    py = torch.diag(u.T @ ry @ u)
    # Pn = diag(E'*Rn*E); %equation (24)
    pn = torch.diag(u.T @ noise_corr @ u)
    # cost_F = -Py + 2 * Pn; % equation(22)
    cost_f = -py + 2 * pn
    # print(cost_f)
    # kf = sum(cost_F < 0);
    sig_subspace_dim = (cost_f < 0).sum()
    # [dummy, ind_asc] = sort(cost_F, 'ascend');
    _, indices = torch.sort(cost_f)
    # Ek = E(:, ind_asc(1: kf));
    eigs_span_subspace = u[:, indices[:sig_subspace_dim]]
    return sig_subspace_dim, eigs_span_subspace
//  [python]
def fit_ellipse(points):
    """
    fit a 2D ellipse in the XY plane for a set of 3D points
    points: instance points Nx3
    return: axis sizes [size_x, size_y, size_z], angle
    """
    N = points.shape[0]
    X = points[:, :2]
    X = X - X.mean(0)
    _out = torch.svd(
        (X - X.mean(axis=0)) / torch.sqrt(torch.tensor(N)), some=False
    )
    angle = np.arctan2(_out[2][0, 1].cpu(), _out[2][0, 0].cpu()).cuda()
    zc = (torch.max(points[:, 2]) + torch.min(points[:, 2])) / 2
    dz = torch.cdist(points[:, 2].unsqueeze(1), zc.unsqueeze(0).unsqueeze(1))
    sz = 2 * dz.mean()
    size = torch.cat((_out[1], torch.tensor([sz]).to(_out[1])))
    return size, angle
//  [python]
def fit_ellipse(points):
    """
    fit a 2D ellipse in the XY plane for a set of 3D points
    points: instance points Nx3
    return: axis sizes [size_x, size_y, size_z], angle
    """
    N = points.shape[0]
    X = points[:, :2]
    X = X - X.mean(0)
    _out = torch.linalg.svd(
        (X - X.mean(axis=0)) / torch.sqrt(torch.tensor(N)), full_matrices=False
    )
    angle = np.arctan2(_out.Vh[0, 1].cpu(), _out.Vh[0, 0].cpu()).cuda()
    zc = (torch.max(points[:, 2]) + torch.min(points[:, 2])) / 2
    dz = torch.cdist(points[:, 2].unsqueeze(1), zc.unsqueeze(0).unsqueeze(1))
    sz = 2 * dz.mean()
    size = torch.cat((_out.S, torch.tensor([sz]).to(_out.S)))
    return size, angle
//  [python]
def test_transfer_parametrizations_and_params_many_to_one(self):
        # A parametrization with several outputs
        class RankOne(nn.Module):
            def forward(self, x, y):
                # Form a rank-1 matrix from a pair of vectors
                return x.unsqueeze(-1) @ y.unsqueeze(-2)

            def right_inverse(self, Y):
                # We project the given matrix onto the rank 1 matrices
                U, S, Vh = torch.svd(Y)
                # S is ordered in a decreasing way.
                s0_sqrt = S[0].sqrt().unsqueeze(-1)
                return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt

        class Double(nn.Module):
            def forward(self, x):
                return 2.0 * x

        model = nn.Linear(3, 3)
        parametrize.register_parametrization(model, "weight", RankOne())
        parametrize.register_parametrization(model, "weight", Double())
        hold_weight = model.weight

        to_model = torch.ao.nn.qat.Linear(
            3, 3, qconfig=torch.ao.quantization.get_default_qconfig()
        )

        parametrize.transfer_parametrizations_and_params(model, to_model)

        # checks that final and original value are correct and the to_model is parametrized
        self.assertTrue(torch.nn.utils.parametrize.is_parametrized(to_model, "weight"))
        self.assertEqual(model.weight, to_model.weight)
        self.assertEqual(
            model.parametrizations.weight.original0,
            to_model.parametrizations.weight.original0,
        )
        self.assertEqual(
            model.parametrizations.weight.original1,
            to_model.parametrizations.weight.original1,
        )

        # check that the transfer didn't affect the original value
        self.assertEqual(hold_weight, model.weight)

        # testing that changes to one set of parametrizations do not affect the other
        model.test_param = Parameter(torch.randn(3, 3))

        self.assertTrue(not hasattr(to_model, "test_param"))
        parametrize.register_parametrization(model, "test_param", RankOne())
        hold_test_param = model.test_param
        parametrize.transfer_parametrizations_and_params(model, to_model, "test_param")

        # also check that previously missing params got transferred correctly
        self.assertEqual(model.test_param, to_model.test_param)
        self.assertEqual(
            model.parametrizations.test_param.original0,
            to_model.parametrizations.test_param.original0,
        )
        self.assertEqual(
            model.parametrizations.test_param.original1,
            to_model.parametrizations.test_param.original1,
        )

        # check that the new transfer didn't change the value for the from_module
        self.assertEqual(hold_test_param, model.test_param)
//  [python]
def test_transfer_parametrizations_and_params_many_to_one(self):
        # A parametrization with several outputs
        class RankOne(nn.Module):
            def forward(self, x, y):
                # Form a rank-1 matrix from a pair of vectors
                return x.unsqueeze(-1) @ y.unsqueeze(-2)

            def right_inverse(self, Y):
                # We project the given matrix onto the rank 1 matrices
                U, S, Vh = torch.linalg.svd(Y, full_matrices=False)
                # S is ordered in a decreasing way.
                s0_sqrt = S[0].sqrt().unsqueeze(-1)
                return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt

        class Double(nn.Module):
            def forward(self, x):
                return 2.0 * x

        model = nn.Linear(3, 3)
        parametrize.register_parametrization(model, "weight", RankOne())
        parametrize.register_parametrization(model, "weight", Double())
        hold_weight = model.weight

        to_model = torch.ao.nn.qat.Linear(
            3, 3, qconfig=torch.ao.quantization.get_default_qconfig()
        )

        parametrize.transfer_parametrizations_and_params(model, to_model)

        # checks that final and original value are correct and the to_model is parametrized
        self.assertTrue(torch.nn.utils.parametrize.is_parametrized(to_model, "weight"))
        self.assertEqual(model.weight, to_model.weight)
        self.assertEqual(
            model.parametrizations.weight.original0,
            to_model.parametrizations.weight.original0,
        )
        self.assertEqual(
            model.parametrizations.weight.original1,
            to_model.parametrizations.weight.original1,
        )

        # check that the transfer didn't affect the original value
        self.assertEqual(hold_weight, model.weight)

        # testing that changes to one set of parametrizations do not affect the other
        model.test_param = Parameter(torch.randn(3, 3))

        self.assertTrue(not hasattr(to_model, "test_param"))
        parametrize.register_parametrization(model, "test_param", RankOne())
        hold_test_param = model.test_param
        parametrize.transfer_parametrizations_and_params(model, to_model, "test_param")

        # also check that previously missing params got transferred correctly
        self.assertEqual(model.test_param, to_model.test_param)
        self.assertEqual(
            model.parametrizations.test_param.original0,
            to_model.parametrizations.test_param.original0,
        )
        self.assertEqual(
            model.parametrizations.test_param.original1,
            to_model.parametrizations.test_param.original1,
        )

        # check that the new transfer didn't change the value for the from_module
        self.assertEqual(hold_test_param, model.test_param)
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    S1, S2 = S1.float(), S2.float()
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2)

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.svd(K)
    V = Vh.transpose(-2,-1)
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    S1, S2 = S1.float(), S2.float()
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2)

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.linalg.svd(K)
    V = Vh.transpose(-2,-1)
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [python]
def _compute_low_rank_preconditioned_gradient(
        self,
        preconditioned_gradient: torch.Tensor,
    ) -> List[torch.Tensor]:
        """Performs low-rank approximation of the preconditioned gradient with SVD.

        Args:
            preconditioned_gradient (torch.Tensor):
                The preconditioned per-sample-gradient matrix to be low-rank approximated.

        Returns:
            List[torch.Tensor, torch.Tensor]:
                Low-rank matrices that approximate the original preconditioned gradient.
        """
        U, S, V = torch.svd(  # pylint: disable=not-callable
            preconditioned_gradient.contiguous().to(dtype=self.score_args.query_gradient_svd_dtype),
            some=False,
        )
        rank = self.score_args.query_gradient_rank
        U_k = U[:, :, :rank]
        S_k = S[:, :rank]
        # Avoid holding the full memory of the original tensor before indexing.
        V_k = V[:, :rank, :].clone()
        return [
            torch.matmul(U_k, torch.diag_embed(S_k)).to(dtype=self.score_args.score_dtype).contiguous(),
            V_k.to(dtype=self.score_args.score_dtype).contiguous(),
        ]
//  [python]
def _compute_low_rank_preconditioned_gradient(
        self,
        preconditioned_gradient: torch.Tensor,
    ) -> List[torch.Tensor]:
        """Performs low-rank approximation of the preconditioned gradient with SVD.

        Args:
            preconditioned_gradient (torch.Tensor):
                The preconditioned per-sample-gradient matrix to be low-rank approximated.

        Returns:
            List[torch.Tensor, torch.Tensor]:
                Low-rank matrices that approximate the original preconditioned gradient.
        """
        U, S, V = torch.linalg.svd(  # pylint: disable=not-callable
            preconditioned_gradient.contiguous().to(dtype=self.score_args.query_gradient_svd_dtype),
            full_matrices=False,
        )
        rank = self.score_args.query_gradient_rank
        U_k = U[:, :, :rank]
        S_k = S[:, :rank]
        # Avoid holding the full memory of the original tensor before indexing.
        V_k = V[:, :rank, :].clone()
        return [
            torch.matmul(U_k, torch.diag_embed(S_k)).to(dtype=self.score_args.score_dtype).contiguous(),
            V_k.to(dtype=self.score_args.score_dtype).contiguous(),
        ]
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,2,1))

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.svd(K)
    V = Vh.transpose(-2,-1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,2,1))

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.linalg.svd(K)
    V = Vh.transpose(-2,-1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def S_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:

                    # rank adaption

                    s_small = torch.clone(l.S_hat[:2 * l.dynamic_rank, :2 * l.dynamic_rank])
                    try:
                        u2, d, v2 = torch.svd(s_small)
                    except Exception as e:
                        print(e)
                        print(s_small)
                        u2, d, v2 = np.linalg.svd(s_small)

                    tmp = 0.0
                    tol = self.theta * torch.norm(d) if not self.absolute else self.theta 
                    rmax = int(np.floor(d.shape[0] / 2))
                    for j in range(0, 2 * rmax - 1):
                        tmp = torch.norm(d[j:2 * rmax - 1])
                        if tmp < tol:
                            rmax = j
                            break

                    rmax = min([rmax, l.rmax])
                    rmax = max([rmax, 2])

                    l.S_hat[:rmax,:rmax] = torch.diag(d[:rmax])
                    l.U[:, :rmax] = l.U_hat[:, :2 * l.dynamic_rank]@u2[:, :rmax]
                    l.V[:,:rmax] =  l.V_hat[:,:2 * l.dynamic_rank]@(v2[:, :rmax])
                    l.dynamic_rank = int(rmax)
//  [python]
def S_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:

                    # rank adaption

                    s_small = torch.clone(l.S_hat[:2 * l.dynamic_rank, :2 * l.dynamic_rank])
                    try:
                        u2, d, v2 = torch.linalg.svd(s_small)
                    except Exception as e:
                        print(e)
                        print(s_small)
                        u2, d, v2 = np.linalg.svd(s_small)

                    tmp = 0.0
                    tol = self.theta * torch.linalg.norm(d) if not self.absolute else self.theta 
                    rmax = int(np.floor(d.shape[0] / 2))
                    for j in range(0, 2 * rmax - 1):
                        tmp = torch.linalg.norm(d[j:2 * rmax - 1])
                        if tmp < tol:
                            rmax = j
                            break

                    rmax = min([rmax, l.rmax])
                    rmax = max([rmax, 2])

                    l.S_hat[:rmax,:rmax] = torch.diag(d[:rmax])
                    l.U[:, :rmax] = l.U_hat[:, :2 * l.dynamic_rank]@u2[:, :rmax]
                    l.V[:,:rmax] =  l.V_hat[:,:2 * l.dynamic_rank]@(v2[:, :rmax])
                    l.dynamic_rank = int(rmax)
//  [python]
def forward(self, data: Data) -> Data:
        assert data.x is not None

        if data.x.size(-1) > self.out_channels:
            U, S, _ = torch.svd(data.x)
            data.x = torch.mm(U[:, :self.out_channels],
                              torch.diag(S[:self.out_channels]))
        return data
//  [python]
def forward(self, data: Data) -> Data:
        assert data.x is not None

        if data.x.size(-1) > self.out_channels:
            U, S, _ = torch.linalg.svd(data.x)
            data.x = torch.mm(U[:, :self.out_channels],
                              torch.diag(S[:self.out_channels]))
        return data
//  [python]
def dual_pca(X:np.ndarray, Y:np.ndarray, 
            dim:Optional[int]=50,
            singular:Optional[bool]=False,
            backend:Optional[str]='sklearn',
            use_gpu:Optional[bool]=True
    ) -> List[Tensor]:
    r"""
    Dual PCA for batch correction
    
    Parameters:
    ----------
    X
        expr matrix 1 in shape of (cells, genes)
    Y
        expr matrix 2 in shape of (cells, genes)
    dim
        dimension of embedding
    singular
        if multiply the singular value
    backend
        backend to calculate singular value
    use_gpu
        if calculate in gpu
    Returns:
    ----------
    embd1, embd2: Tensors of embedding
    
    References:
    ----------
    Thanks Xin-Ming Tu for his [blog](https://xinmingtu.cn/blog/2022/CCA_dual_PCA/)
    """
    assert X.shape[1] == Y.shape[1]
    device = torch.device(f'cuda:{get_free_gpu()}' if torch.cuda.is_available() and use_gpu else 'cpu')
    X = torch.Tensor(X).to(device=device)
    Y = torch.Tensor(Y).to(device=device)
    cor_var = X @ Y.T
    if backend == 'torch':
        U, S, Vh = torch.svd(cor_var)
        if not singular:
            return U[:, :dim], Vh.T[:, :dim]
        Z_x = U[:, :dim] @ torch.sqrt(torch.diag(S[:dim]))
        Z_y = Vh.T[:, :dim] @ torch.sqrt(torch.diag(S[:dim]))
        return Z_x.cpu(), Z_y.cpu()
        # torch.dist(cor_var, Z_x @ Z_y.T)  # check the information loss
    elif backend == 'sklearn':
        cor_var = cor_var.cpu().numpy()
        U, S, Vh = randomized_svd(cor_var, n_components=dim, random_state=0)
        if not singular:
            return Tensor(U), Tensor(Vh.T)
        Z_x = U @ np.sqrt(np.diag(S))
        Z_y = Vh.T @ np.sqrt(np.diag(S))
        return Tensor(Z_x), Tensor(Z_y)
//  [python]
def dual_pca(X:np.ndarray, Y:np.ndarray, 
            dim:Optional[int]=50,
            singular:Optional[bool]=False,
            backend:Optional[str]='sklearn',
            use_gpu:Optional[bool]=True
    ) -> List[Tensor]:
    r"""
    Dual PCA for batch correction
    
    Parameters:
    ----------
    X
        expr matrix 1 in shape of (cells, genes)
    Y
        expr matrix 2 in shape of (cells, genes)
    dim
        dimension of embedding
    singular
        if multiply the singular value
    backend
        backend to calculate singular value
    use_gpu
        if calculate in gpu
    Returns:
    ----------
    embd1, embd2: Tensors of embedding
    
    References:
    ----------
    Thanks Xin-Ming Tu for his [blog](https://xinmingtu.cn/blog/2022/CCA_dual_PCA/)
    """
    assert X.shape[1] == Y.shape[1]
    device = torch.device(f'cuda:{get_free_gpu()}' if torch.cuda.is_available() and use_gpu else 'cpu')
    X = torch.Tensor(X).to(device=device)
    Y = torch.Tensor(Y).to(device=device)
    cor_var = X @ Y.T
    if backend == 'torch':
        U, S, Vh = torch.linalg.svd(cor_var)
        if not singular:
            return U[:, :dim], Vh.T[:, :dim]
        Z_x = U[:, :dim] @ torch.sqrt(torch.diag(S[:dim]))
        Z_y = Vh.T[:, :dim] @ torch.sqrt(torch.diag(S[:dim]))
        return Z_x.cpu(), Z_y.cpu()
        # torch.dist(cor_var, Z_x @ Z_y.T)  # check the information loss
    elif backend == 'sklearn':
        cor_var = cor_var.cpu().numpy()
        U, S, Vh = randomized_svd(cor_var, n_components=dim, random_state=0)
        if not singular:
            return Tensor(U), Tensor(Vh.T)
        Z_x = U @ np.sqrt(np.diag(S))
        Z_y = Vh.T @ np.sqrt(np.diag(S))
        return Tensor(Z_x), Tensor(Z_y)
//  [python]
def fit_line(points: Tensor, weights: Optional[Tensor] = None) -> ParametrizedLine:
    """Fit a line from a set of points.

    Args:
        points: tensor containing a batch of sets of n-dimensional points. The  expected
            shape of the tensor is :math:`(B,N,D)`.
        weights: weights to use to solve the equations system. The  expected
            shape of the tensor is :math:`(B,N)`.

    Return:
        A tensor containing the direction of the fited line of shape :math:`(B,D)`.

    Example:
        >>> points = torch.rand(2, 10, 3)
        >>> weights = torch.ones(2, 10)
        >>> line = fit_line(points, weights)
        >>> line.direction.shape
        torch.Size([2, 3])
    """
    KORNIA_CHECK_IS_TENSOR(points, "points must be a tensor")
    KORNIA_CHECK_SHAPE(points, ["B", "N", "D"])

    mean = points.mean(-2, True)
    A = points - mean

    if weights is not None:
        KORNIA_CHECK_IS_TENSOR(weights, "weights must be a tensor")
        KORNIA_CHECK_SHAPE(weights, ["B"])
        KORNIA_CHECK(points.shape[0] == weights.shape[0])
        A = A.transpose(-2, -1) @ torch.diag_embed(weights) @ A
    else:
        A = A.transpose(-2, -1) @ A

    # NOTE: not optimal for 2d points, but for now works for other dimensions
    U, S, V = torch.svd(A)

    # the first left eigenvector is the direction on the fited line
    direction = V[..., 0, :]  # BxD
    origin = mean[..., 0, :]  # BxD

    return ParametrizedLine(origin, direction)
//  [python]
def fit_line(points: Tensor, weights: Optional[Tensor] = None) -> ParametrizedLine:
    """Fit a line from a set of points.

    Args:
        points: tensor containing a batch of sets of n-dimensional points. The  expected
            shape of the tensor is :math:`(B,N,D)`.
        weights: weights to use to solve the equations system. The  expected
            shape of the tensor is :math:`(B,N)`.

    Return:
        A tensor containing the direction of the fited line of shape :math:`(B,D)`.

    Example:
        >>> points = torch.rand(2, 10, 3)
        >>> weights = torch.ones(2, 10)
        >>> line = fit_line(points, weights)
        >>> line.direction.shape
        torch.Size([2, 3])
    """
    KORNIA_CHECK_IS_TENSOR(points, "points must be a tensor")
    KORNIA_CHECK_SHAPE(points, ["B", "N", "D"])

    mean = points.mean(-2, True)
    A = points - mean

    if weights is not None:
        KORNIA_CHECK_IS_TENSOR(weights, "weights must be a tensor")
        KORNIA_CHECK_SHAPE(weights, ["B"])
        KORNIA_CHECK(points.shape[0] == weights.shape[0])
        A = A.transpose(-2, -1) @ torch.diag_embed(weights) @ A
    else:
        A = A.transpose(-2, -1) @ A

    # NOTE: not optimal for 2d points, but for now works for other dimensions
    _, _, V = torch.linalg.svd(A)

    # the first left eigenvector is the direction on the fited line
    direction = V[..., 0, :]  # BxD
    origin = mean[..., 0, :]  # BxD

    return ParametrizedLine(origin, direction)
//  [python]
def _svd_lowrank(
    A: Tensor,
    q: Optional[int] = 6,
    niter: Optional[int] = 2,
    M: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t, some=False)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t, some=False)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def _svd_lowrank(
    A: Tensor,
    q: Optional[int] = 6,
    niter: Optional[int] = 2,
    M: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def compute_power_svd(matrix: torch.Tensor, power: float) -> torch.Tensor:
    r"""Compute G^{-1/p} using a SVD.

        Calculate SVD on the GPU. Sometimes, SVD on the CPU is faster than GPU, but based on the several experiments,
        CUDA seems much faster than on CPU.

    :param matrix: torch.Tensor. a square positive semi-definite matrix.
    :param power: float. rank.
    """
    u, s, vh = torch.svd(matrix)
    s.pow_(-1.0 / power)
    return u @ (s.diag() if len(matrix.shape) == 2 else s.diag_embed()) @ vh
//  [python]
def compute_power_svd(matrix: torch.Tensor, power: float) -> torch.Tensor:
    r"""Compute G^{-1/p} using a SVD.

        Calculate SVD on the GPU. Sometimes, SVD on the CPU is faster than GPU, but based on the several experiments,
        CUDA seems much faster than on CPU.

    :param matrix: torch.Tensor. a square positive semi-definite matrix.
    :param power: float. rank.
    """
    u, s, vh = torch.linalg.svd(matrix, full_matrices=False)
    s.pow_(-1.0 / power)
    return u @ (s.diag() if len(matrix.shape) == 2 else s.diag_embed()) @ vh
//  [python]
def init_text_embeddings(self):

        label2text_dict = self.label2text_dict
        transformed_dict = {key: value.replace('-', ' ').lower() for key, value in label2text_dict.items()}

        clip_dict = {key: f"A magnetic resonance imaging of a {value} in human brain." for key, value in transformed_dict.items()}

        text_segments = clip.tokenize(list(clip_dict.values())).to('cuda')
        model, _ = clip.load(self.clip_backbone, 'cuda')
        model.eval()
        model.requires_grad_(False)
        text_features = model.encode_text(text_segments).float() # 35 x 512
        U, S, V = torch.svd(text_features)
        background_vector = V[:, -1].unsqueeze(0)
        self.text_features = torch.cat([background_vector, text_features], dim=0) # 36 x 512
        self.text_features.requires_grad_(False)
        self.text_features = self.text_features.float()
        print("Successfully initialized *CLIP* embedded anatomical features for %d classes using %s backbone." % (self.text_features.shape[0], self.clip_backbone))

        self.init_dim = self.text_features.shape[1]
//  [python]
def init_text_embeddings(self):

        label2text_dict = self.label2text_dict
        transformed_dict = {key: value.replace('-', ' ').lower() for key, value in label2text_dict.items()}

        clip_dict = {key: f"A magnetic resonance imaging of a {value} in human brain." for key, value in transformed_dict.items()}

        text_segments = clip.tokenize(list(clip_dict.values())).to('cuda')
        model, _ = clip.load(self.clip_backbone, 'cuda')
        model.eval()
        model.requires_grad_(False)
        text_features = model.encode_text(text_segments).float() # 35 x 512
        _, _, V = torch.linalg.svd(text_features)
        background_vector = V[:, -1].unsqueeze(0)
        self.text_features = torch.cat([background_vector, text_features], dim=0) # 36 x 512
        self.text_features.requires_grad_(False)
        self.text_features = self.text_features.float()
        print("Successfully initialized *CLIP* embedded anatomical features for %d classes using %s backbone." % (self.text_features.shape[0], self.clip_backbone))

        self.init_dim = self.text_features.shape[1]
//  [python]
def forward(self, query_feat, support_feat, way_num, shot_num, normalize=True, discriminative=False):
        e, ws, d = support_feat.size()
        support_feat = support_feat.reshape(e, way_num, shot_num, -1)
        query_feat = query_feat.unsqueeze(1)
        try:
            UU, _, _ = torch.svd(support_feat.permute(0, 1, 3, 2).double())
        except AttributeError:
            UU, _, _ = torch.svd(support_feat.permute(0, 1, 3, 2).double())
        UU = UU.float()
        subspace = UU[:, :, :, :shot_num - 1].permute(0, 1, 3, 2)
        projection = subspace.permute(0, 1, 3, 2).matmul(subspace.matmul(query_feat.permute(0, 1, 3, 2))).permute(0, 1, 3, 2)
        dist = torch.sum((query_feat - projection) ** 2, dim=-1).permute(0, 2, 1)
        logits = -dist
        if normalize:
            logits /= d
        disc_loss = 0
        if discriminative:
            subspace_metric = torch.norm(torch.matmul(subspace.unsqueeze(1), subspace.unsqueeze(2).transpose(-2, -1)), p='fro', dim=[-2, -1])
            mask = torch.eye(way_num).bool()
            subspace_metric = subspace_metric[:, ~mask]
            disc_loss = torch.sum(subspace_metric ** 2)
        return logits, disc_loss
//  [python]
def forward(self, query_feat, support_feat, way_num, shot_num, normalize=True, discriminative=False):
        e, ws, d = support_feat.size()
        support_feat = support_feat.reshape(e, way_num, shot_num, -1)
        query_feat = query_feat.unsqueeze(1)
        try:
            UU, _, _ = torch.linalg.svd(support_feat.permute(0, 1, 3, 2).double())
        except AttributeError:
            UU, _, _ = torch.svd(support_feat.permute(0, 1, 3, 2).double())
        UU = UU.float()
        subspace = UU[:, :, :, :shot_num - 1].permute(0, 1, 3, 2)
        projection = subspace.permute(0, 1, 3, 2).matmul(subspace.matmul(query_feat.permute(0, 1, 3, 2))).permute(0, 1, 3, 2)
        dist = torch.sum((query_feat - projection) ** 2, dim=-1).permute(0, 2, 1)
        logits = -dist
        if normalize:
            logits /= d
        disc_loss = 0
        if discriminative:
            subspace_metric = torch.norm(torch.matmul(subspace.unsqueeze(1), subspace.unsqueeze(2).transpose(-2, -1)), p='fro', dim=[-2, -1])
            mask = torch.eye(way_num).bool()
            subspace_metric = subspace_metric[:, ~mask]
            disc_loss = torch.sum(subspace_metric ** 2)
        return logits, disc_loss
//  [python]
def fit(data, models_posterior):
    # pts: B x C x H x W / B x C x N
    # models_likelihood: B x K x H x W
    # means, stds: B x K x C
    # out: argmax(P(H|E)) B x K x C

    pts = data["pts"]
    B, K, H, W = models_posterior.shape
    # models_posterior = (models_posterior > 0.5) * 1.0
    # B x K x 3
    centroid = (models_posterior[:, :, None] * pts[:, None]).flatten(3).sum(dim=3) / (
        models_posterior[:, :, None].flatten(3).sum(dim=3) + 1e-8
    )
    # B x K x 3 x N
    dev = (
        models_posterior[:, :, None] * (pts[:, None] - centroid[:, :, :, None, None])
    ).flatten(3)
    dev = dev.permute(0, 1, 3, 2)
    # B x K x N x 3

    U, stds, axesT = torch.svd(dev)
    # B x K x 3
    axes = axesT.permute(0, 1, 3, 2)

    # B x K x N x 3
    dev_proj = dev @ axes
    # B x K x 3
    dev_max, _ = dev_proj.abs().max(dim=2)
    measures = dev_max * 2
    params = {}
    params["center"] = centroid
    params["axes"] = axes
    params["measures"] = measures
    params["stds"] = stds
    params["pts"] = []

    for b in range(B):
        params["pts"].append([])
        for k in range(K):
            params["pts"][b].append(pts[b, :, models_posterior[b, k] > 0.5])

    o4visual.visualize_boxes3d(params)
    # +- 1 std: 68.27%
    # +- 2 std: 95.45%
    # +- 3 std: 99.73%
    return params