jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _fit(self, embeddings: torch.Tensor) -> None:
        """
        Perform the PCA. Evaluate ``components``, ``expoained_variance``, ``explained_variance_ratio``,
        ``singular_values``, ``mean``.

        Args:
            embeddings: Embeddings matrix with the shape of ``[n_embeddings, embeddings_dim]``.
        """
        # The code below mirrors code from scikit-learn repository:
        # https://github.com/scikit-learn/scikit-learn/blob/f3f51f9b6/sklearn/decomposition/_pca.py#L491
        n_samples = embeddings.shape[0]
        self.mean = embeddings.mean(dim=0).unsqueeze(0)
        embeddings = embeddings - self.mean
        # if there are more embeddings than its dimension, then we will not perform full matrices evaluation
        full_matrices = embeddings.shape[0] < embeddings.shape[1]
        _, self.singular_values, self.components = torch.linalg.svd(embeddings, full_matrices=full_matrices)
        self.explained_variance = self.singular_values**2 / (n_samples - 1)
        self.explained_variance_ratio = self.explained_variance / self.explained_variance.sum()

        # Make components deterministic.
        # Note. In sklearn this operation is done based on the U matrix of SVD decomposition, and
        # here V matrix is used. So, the components of this class and sklearn.decomposition.PCA could differ in sign.
        # See the following for details:
        # https://github.com/scikit-learn/scikit-learn/blob/f3f51f9b6/sklearn/decomposition/_pca.py#L520
        n_components = self.components.shape[0]
        max_abs_rows = torch.argmax(torch.abs(self.components), dim=1)
        signs = torch.sign(self.components[torch.arange(n_components), max_abs_rows])
        self.components *= signs.unsqueeze(1)
//  [python]
def decode_E(E):
    U, S, Vt = torch.svd(E)
    S_diag = torch.zeros_like(U)
    S_diag[0, 0] = (S[0] + S[1])/2.0
    S_diag[1, 1] = (S[0] + S[1])/2.0
    W = torch.tensor([[0.0, -1., 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
    Z = torch.tensor([[0.0, 1.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
    rec_R1 = U @ W @ Vt
    if torch.det(rec_R1) < 0:
        rec_R1 = -rec_R1

    rec_R2 = U @ W.T @ Vt
    if torch.det(rec_R2) < 0:
        rec_R2 = -rec_R2
    
    rec_T = U @ W @ S_diag @ U.T
    rec_t1 = skew2vec(rec_T)
    
    return rec_R1, rec_R2, rec_t1, -rec_t1
//  [python]
def decode_E(E):
    U, S, Vt = torch.linalg.svd(E)
    S_diag = torch.zeros_like(U)
    S_diag[0, 0] = (S[0] + S[1])/2.0
    S_diag[1, 1] = (S[0] + S[1])/2.0
    W = torch.tensor([[0.0, -1., 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
    Z = torch.tensor([[0.0, 1.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
    rec_R1 = U @ W @ Vt
    if torch.linalg.det(rec_R1) < 0:
        rec_R1 = -rec_R1

    rec_R2 = U @ W.T @ Vt
    if torch.linalg.det(rec_R2) < 0:
        rec_R2 = -rec_R2
    
    rec_T = U @ W @ S_diag @ U.T
    rec_t1 = skew2vec(rec_T)
    
    return rec_R1, rec_R2, rec_t1, -rec_t1
//  [python]
def forward(self, A):
        U, S, V = torch.svd(A)
        #U, S, Vt = torch.linalg.svd(A, full_matrices=False)
        #V = Vt.t()
        self.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def forward(self, A):
        #U, S, V = torch.svd(A)
        U, S, Vt = torch.linalg.svd(A, full_matrices=False)
        V = Vt.t()
        self.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def svd_approx(
    W: torch.Tensor, r: int | None = None
) -> tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]:
    """
    Perform SVD approximation of a matrix with an optional rank specification.

    Parameters:
    W (torch.Tensor): The input matrix to be approximated.
    r (int, optional): The rank for the approximation. If None, the full rank is used.

    Returns:
    Tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]:
        A tuple containing the matrices A and B (where A @ B is the approximation of W),
        the rank used for the approximation, and the Frobenius norm of the difference between W and its approximation.
    """
    U, S, V = torch.svd(W)
    if r is None:
        r = S.size(0)
    else:
        r = min(r, S.size(0))

    A = U[:, :r]
    B = torch.diag(S[:r]) @ V[:r, :]
    approx_W = A @ B
    frobenius_norm = torch.linalg.norm(W - approx_W)
    return A, B, r, frobenius_norm
//  [python]
def svd_approx(
    W: torch.Tensor, r: int | None = None
) -> tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]:
    """
    Perform SVD approximation of a matrix with an optional rank specification.

    Parameters:
    W (torch.Tensor): The input matrix to be approximated.
    r (int, optional): The rank for the approximation. If None, the full rank is used.

    Returns:
    Tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]:
        A tuple containing the matrices A and B (where A @ B is the approximation of W),
        the rank used for the approximation, and the Frobenius norm of the difference between W and its approximation.
    """
    U, S, V = torch.linalg.svd(W, full_matrices=False)
    if r is None:
        r = S.size(0)
    else:
        r = min(r, S.size(0))

    A = U[:, :r]
    B = torch.diag(S[:r]) @ V[:r, :]
    approx_W = A @ B
    frobenius_norm = torch.linalg.norm(W - approx_W)
    return A, B, r, frobenius_norm
//  [python]
def compute_decoder_jacobian_singular_values(model, z, conditions=None):
    """
    Compute the singular values of the Jacobian of the decoder with respect to the input z.
    Args:
        model (nn.Module): A model with implemented decode() method.
        z (Tensor): The input to the decode() method.
        conditions (Tensor): The conditions to the decode() method.
    Returns:
        Tensor: The singular values of the Jacobian of the decoder with respect to the input z.
    """
    jacobians = compute_decoder_jacobian(model, z, conditions)
    return torch.vmap(
        lambda jacobian: torch.svd(jacobian, compute_uv=False)
    )(jacobians)
//  [python]
def compute_decoder_jacobian_singular_values(model, z, conditions=None):
    """
    Compute the singular values of the Jacobian of the decoder with respect to the input z.
    Args:
        model (nn.Module): A model with implemented decode() method.
        z (Tensor): The input to the decode() method.
        conditions (Tensor): The conditions to the decode() method.
    Returns:
        Tensor: The singular values of the Jacobian of the decoder with respect to the input z.
    """
    jacobians = compute_decoder_jacobian(model, z, conditions)
    return torch.vmap(
        lambda jacobian: torch.linalg.svd(jacobian, full_matrices=False)[1]
    )(jacobians)
//  [python]
def svd_distill(residual, rank, clamp_quantile=1):
    residual = residual.float()
    U, S, Vh = torch.svd(residual)
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    return U, Vh
//  [python]
def svd_distill(residual, rank, clamp_quantile=1):
    residual = residual.float()
    U, S, Vh = torch.linalg.svd(residual)
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    return U, Vh
//  [python]
def nearest_rotation(mat):
    shape = mat.shape
    mat = mat.reshape(-1, 3, 3)
    U, s, V = torch.svd(mat)
    res = U @ V
    det = torch.det(res)
    mask = det < 0
    if mask.any():
        U[mask, -1, :] = -U[mask, -1, :]
        res = U @ V
    assert (torch.allclose(torch.det(res), torch.ones_like(det)))
    return res.reshape(shape)
//  [python]
def nearest_rotation(mat):
    shape = mat.shape
    mat = mat.reshape(-1, 3, 3)
    U, s, V = torch.linalg.svd(mat)
    res = U @ V
    det = torch.linalg.det(res)
    mask = det < 0
    if mask.any():
        U[mask, -1, :] = -U[mask, -1, :]
        res = U @ V
    assert (torch.allclose(torch.linalg.det(res), torch.ones_like(det)))
    return res.reshape(shape)
//  [python]
def morlets(grid_or_shape, theta, xis, sigmas, slants, morlet=True, ifftshift=True, fft=True):
    """Creates morlet wavelet filters from inputs

        Parameters:
            grid_or_shape -- a grid of the size of the filter or a tuple that indicates its shape
            theta -- global orientations of the wavelets
            xis -- frequency scales of the wavelets
            sigmas -- gaussian window scales of the wavelets
            slants -- slants of the wavelets
            morlet -- boolean for morlet or gabor wavelet
            ifftshift -- boolean for the ifftshift (inverse fast fourier transform shift)
            fft -- boolean for the fft (fast fourier transform)

        Returns:
            wavelets -- the wavelet filters

    """
    orientations = torch.cat((torch.cos(theta).unsqueeze(1),torch.sin(theta).unsqueeze(1)), dim =1)
    n_filters, ndim = orientations.shape
    wave_vectors = orientations * xis[:, np.newaxis]
    _, _, gauss_directions = torch.svd(orientations[:, np.newaxis])
    gauss_directions = gauss_directions / sigmas[:, np.newaxis, np.newaxis]
    indicator = torch.arange(ndim, device=slants.device) < 1
    slant_modifications = (1.0 * indicator + slants[:, np.newaxis] * ~indicator)
    gauss_directions = gauss_directions * slant_modifications[:, :, np.newaxis]

    wavelets = raw_morlets(grid_or_shape, wave_vectors, gauss_directions, morlet=morlet, 
                          ifftshift=ifftshift, fft=fft)

    norm_factors = (2 * 3.1415 * sigmas * sigmas / slants).unsqueeze(1)

    if type(grid_or_shape) == tuple:
        norm_factors = norm_factors.expand([n_filters,grid_or_shape[0]]).unsqueeze(2).repeat(1,1,grid_or_shape[1])
    else:
        norm_factors = norm_factors.expand([n_filters,grid_or_shape.shape[1]]).unsqueeze(2).repeat(1,1,grid_or_shape.shape[2])

    wavelets = wavelets / norm_factors

    return wavelets
//  [python]
def morlets(grid_or_shape, theta, xis, sigmas, slants, morlet=True, ifftshift=True, fft=True):
    """Creates morlet wavelet filters from inputs

        Parameters:
            grid_or_shape -- a grid of the size of the filter or a tuple that indicates its shape
            theta -- global orientations of the wavelets
            xis -- frequency scales of the wavelets
            sigmas -- gaussian window scales of the wavelets
            slants -- slants of the wavelets
            morlet -- boolean for morlet or gabor wavelet
            ifftshift -- boolean for the ifftshift (inverse fast fourier transform shift)
            fft -- boolean for the fft (fast fourier transform)

        Returns:
            wavelets -- the wavelet filters

    """
    orientations = torch.cat((torch.cos(theta).unsqueeze(1),torch.sin(theta).unsqueeze(1)), dim =1)
    n_filters, ndim = orientations.shape
    wave_vectors = orientations * xis[:, np.newaxis]
    _, _, gauss_directions = torch.linalg.svd(orientations[:, np.newaxis])
    gauss_directions = gauss_directions / sigmas[:, np.newaxis, np.newaxis]
    indicator = torch.arange(ndim, device=slants.device) < 1
    slant_modifications = (1.0 * indicator + slants[:, np.newaxis] * ~indicator)
    gauss_directions = gauss_directions * slant_modifications[:, :, np.newaxis]

    wavelets = raw_morlets(grid_or_shape, wave_vectors, gauss_directions, morlet=morlet, 
                          ifftshift=ifftshift, fft=fft)

    norm_factors = (2 * 3.1415 * sigmas * sigmas / slants).unsqueeze(1)

    if type(grid_or_shape) == tuple:
        norm_factors = norm_factors.expand([n_filters,grid_or_shape[0]]).unsqueeze(2).repeat(1,1,grid_or_shape[1])
    else:
        norm_factors = norm_factors.expand([n_filters,grid_or_shape.shape[1]]).unsqueeze(2).repeat(1,1,grid_or_shape.shape[2])

    wavelets = wavelets / norm_factors

    return wavelets
//  [python]
def pca(X, k):
    """Dimensionality reduction via Principal component analysis / truncated SVD.
    Follows the sklearn implementation:
    https://github.com/scikit-learn/scikit-learn/blob/baf0ea25d6dd034403370fea552b21a6776bef18/sklearn/decomposition/_pca.py#L591
    """
    # center data
    X -= X.mean(dim=-2, keepdim=True)
    # singular value decomposition
    U, S, Vt = torch.svd(X)
    # truncate
    # one way to dim-reduce is to project original input X onto
    # the principal components
    # components = Vt[..., :k, :]
    # i.e., X @ components.transpose(-1, -2)
    # but sklearn's approach is probably more efficient
    # https://github.com/scikit-learn/scikit-learn/blob/baf0ea25d6dd034403370fea552b21a6776bef18/sklearn/decomposition/_pca.py#L434
    return U[..., :k] * S.unsqueeze(-2)[..., :k]
//  [python]
def pca(X, k):
    """Dimensionality reduction via Principal component analysis / truncated SVD.
    Follows the sklearn implementation:
    https://github.com/scikit-learn/scikit-learn/blob/baf0ea25d6dd034403370fea552b21a6776bef18/sklearn/decomposition/_pca.py#L591
    """
    # center data
    X -= X.mean(dim=-2, keepdim=True)
    # singular value decomposition
    U, S, Vt = torch.linalg.svd(X)
    # truncate
    # one way to dim-reduce is to project original input X onto
    # the principal components
    # components = Vt[..., :k, :]
    # i.e., X @ components.transpose(-1, -2)
    # but sklearn's approach is probably more efficient
    # https://github.com/scikit-learn/scikit-learn/blob/baf0ea25d6dd034403370fea552b21a6776bef18/sklearn/decomposition/_pca.py#L434
    return U[..., :k] * S.unsqueeze(-2)[..., :k]
//  [python]
def calc_up_down(mat, dim, conv_dim, device="cpu", clamp_quantile=0.99):
    # if conv_dim is None, diffs do not include LoRAs for conv2d-3x3
    mat = mat.float()
    conv2d = len(mat.size()) == 4
    kernel_size = None if not conv2d else mat.size()[2:4]
    conv2d_3x3 = conv2d and kernel_size != (1, 1)

    rank = dim if not conv2d_3x3 or conv_dim is None else conv_dim
    out_dim, in_dim = mat.size()[0:2]

    if device:
        mat = mat.to(device)

    # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
    rank = min(rank, in_dim, out_dim)  # LoRA rank cannot exceed the original dim

    if conv2d:
        if conv2d_3x3:
            mat = mat.flatten(start_dim=1)
        else:
            mat = mat.squeeze()

    U, S, Vh = torch.svd(mat)

    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if conv2d:
        U = U.reshape(out_dim, rank, 1, 1)
        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])

    U = U.detach().cpu().contiguous()
    Vh = Vh.detach().cpu().contiguous()

    return U, Vh
//  [python]
def calc_up_down(mat, dim, conv_dim, device="cpu", clamp_quantile=0.99):
    # if conv_dim is None, diffs do not include LoRAs for conv2d-3x3
    mat = mat.float()
    conv2d = len(mat.size()) == 4
    kernel_size = None if not conv2d else mat.size()[2:4]
    conv2d_3x3 = conv2d and kernel_size != (1, 1)

    rank = dim if not conv2d_3x3 or conv_dim is None else conv_dim
    out_dim, in_dim = mat.size()[0:2]

    if device:
        mat = mat.to(device)

    # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
    rank = min(rank, in_dim, out_dim)  # LoRA rank cannot exceed the original dim

    if conv2d:
        if conv2d_3x3:
            mat = mat.flatten(start_dim=1)
        else:
            mat = mat.squeeze()

    U, S, Vh = torch.linalg.svd(mat)

    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if conv2d:
        U = U.reshape(out_dim, rank, 1, 1)
        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])

    U = U.detach().cpu().contiguous()
    Vh = Vh.detach().cpu().contiguous()

    return U, Vh
//  [python]
def test_spectral_computation(N: int):
    # Create synthetic data with non-trivial covariance and 0 mean
    # torch.manual_seed(16380)
    mean = torch.tensor([0.0, 0.0, 0.0])
    cov = torch.tensor([[2.0, 1.0, 0.5], [1.0, 4.0, 0.0], [0.5, 0.0, 3.0]])
    dist = torch.distributions.MultivariateNormal(mean, cov)
    data = dist.sample((N,))
    data = data - data.mean(dim=0)

    # Compute svd
    usvh = torch.svd(data, some=False)
    assert usvh[1][0] == usvh[1].max()  # descending order
    v_direct = usvh[2][0, :]
    assert v_direct.numel() == data.size(1)
    assert v_direct.ndim == 1
    assert torch.allclose(usvh[2].square().sum(dim=0), torch.Tensor([1.0]))

    # Compute singular vectors via covariance matrix
    eigs = torch.linalg.eigh(torch.cov(data.T))
    assert eigs.eigenvalues[-1] == eigs.eigenvalues.max()  # ascending order
    v_indirect = eigs.eigenvectors[:, -1]

    # Check that variables are what we think they are
    assert torch.allclose(eigs.eigenvectors.square().sum(dim=1), torch.Tensor([1.0]))
    assert torch.allclose(eigs.eigenvectors.flip(-1).abs(), usvh[2].mH.abs())
    assert torch.allclose(
        eigs.eigenvalues.flip(0), usvh[1].square() / (data.size(0) - 1)
    )

    # Check what we actually care about
    try:
        assert torch.allclose(v_direct, v_indirect)
    except AssertionError:
        # Sign ambiguity
        assert torch.allclose(-v_direct, v_indirect)
//  [python]
def test_spectral_computation(N: int):
    # Create synthetic data with non-trivial covariance and 0 mean
    # torch.manual_seed(16380)
    mean = torch.tensor([0.0, 0.0, 0.0])
    cov = torch.tensor([[2.0, 1.0, 0.5], [1.0, 4.0, 0.0], [0.5, 0.0, 3.0]])
    dist = torch.distributions.MultivariateNormal(mean, cov)
    data = dist.sample((N,))
    data = data - data.mean(dim=0)

    # Compute svd
    usvh = torch.linalg.svd(data, full_matrices=False)
    assert usvh.S[0] == usvh.S.max()  # descending order
    v_direct = usvh.Vh[0, :]
    assert v_direct.numel() == data.size(1)
    assert v_direct.ndim == 1
    assert torch.allclose(usvh.Vh.square().sum(dim=0), torch.Tensor([1.0]))

    # Compute singular vectors via covariance matrix
    eigs = torch.linalg.eigh(torch.cov(data.T))
    assert eigs.eigenvalues[-1] == eigs.eigenvalues.max()  # ascending order
    v_indirect = eigs.eigenvectors[:, -1]

    # Check that variables are what we think they are
    assert torch.allclose(eigs.eigenvectors.square().sum(dim=1), torch.Tensor([1.0]))
    assert torch.allclose(eigs.eigenvectors.flip(-1).abs(), usvh.Vh.mH.abs())
    assert torch.allclose(
        eigs.eigenvalues.flip(0), usvh.S.square() / (data.size(0) - 1)
    )

    # Check what we actually care about
    try:
        assert torch.allclose(v_direct, v_indirect)
    except AssertionError:
        # Sign ambiguity
        assert torch.allclose(-v_direct, v_indirect)
//  [python]
def fit_transform(self, X, y=None, init=None):
        """Fit and transform each point into coordinate space of top n_components PCs.

        To convert back from PC-coordinates to points on the manifold, use ManifoldPCA.inverse_transform.

        :param X: Iterable of points in the space
        :param y: unused
        :param init: unused
        :return: n by n_components matrix of coordinates
        """
        points = self._validate_data(X)

        if y is not None or init is not None:
            warnings.warn("ManifoldPCA does not use 'y' nor 'init' arguments")

        # Estimate the Frechet mean, using the iterative mean method for initialization
        self.frechet_mean_ = optimize_frechet_mean(self.space, points, init_method="iterative")

        # Get the coordinates of every point in the tangent space of the mean
        tangent_vectors = torch.stack([self.space.log_map(self.frechet_mean_, pt) for pt in points],
                                      dim=0)

        # Now we do PCA in the linear-looking tangent space
        _, s, vT = torch.svd(tangent_vectors.view(len(tangent_vectors), -1))
        self.singular_values_ = s
        self.components_ = vT[:self.n_components, :].T.reshape(self.space.shape + (self.n_components,))
        self.scales_ = torch.sqrt(s[:self.n_components])

        return self._transform(tangent_vectors)
//  [python]
def fit_transform(self, X, y=None, init=None):
        """Fit and transform each point into coordinate space of top n_components PCs.

        To convert back from PC-coordinates to points on the manifold, use ManifoldPCA.inverse_transform.

        :param X: Iterable of points in the space
        :param y: unused
        :param init: unused
        :return: n by n_components matrix of coordinates
        """
        points = self._validate_data(X)

        if y is not None or init is not None:
            warnings.warn("ManifoldPCA does not use 'y' nor 'init' arguments")

        # Estimate the Frechet mean, using the iterative mean method for initialization
        self.frechet_mean_ = optimize_frechet_mean(self.space, points, init_method="iterative")

        # Get the coordinates of every point in the tangent space of the mean
        tangent_vectors = torch.stack([self.space.log_map(self.frechet_mean_, pt) for pt in points],
                                      dim=0)

        # Now we do PCA in the linear-looking tangent space
        _, s, vT = torch.linalg.svd(tangent_vectors.view(len(tangent_vectors), -1), full_matrices=False)
        self.singular_values_ = s
        self.components_ = vT[:self.n_components, :].T.reshape(self.space.shape + (self.n_components,))
        self.scales_ = torch.sqrt(s[:self.n_components])

        return self._transform(tangent_vectors)
//  [python]
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)

        loss.backward()
        #Apply Nearest Orthogonal Gradient Mapping
        if args.nog:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad
            else:
                weight_grad = net.module.conv1.weight.grad
            B,C,H,W = weight_grad.size()
            u,s,vh=torch.svd(weight_grad.view(B,C*H*W))
            s[:] = 1
            if args.ow:
                net.module.conv1.random_conv_filter.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
            else:
                net.module.conv1.weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
        #Apply Optimal Learnig Rate
        if args.olr:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad.view(-1, 1)
                weight = net.module.conv1.random_conv_filter.view(-1,1)
            else:
                weight_grad = net.module.conv1.weight.grad.view(-1, 1)
                weight = net.module.conv1.weight.view(-1, 1)
            wtw = weight.t().mm(weight)
            ltw=weight_grad.t().mm(weight)
            ltl=weight_grad.t().mm(weight_grad)
            lr = (wtw*ltw)/(wtw*ltl+2*ltw*ltw)
            if lr>0 and lr < optimizer.param_groups[1]['lr']:
                optimizer.param_groups[0]['lr'] = lr
            else:
                optimizer.param_groups[0]['lr'] = optimizer.param_groups[1]['lr']
        optimizer.step()


        if batch_idx % 500 == 0:
            writer.add_scalar('loss/train_loss', loss.item(), epoch * len(trainloader) + batch_idx + 1)
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    acc = 100. * correct / total
    writer.add_scalar('loss/train_loss_avg', train_loss / len(trainloader), epoch)
    writer.add_scalar('train/accuracy', acc, epoch)
    writer.add_scalar('train/error', 100 - acc, epoch)
    print('current training error {}'.format(100 - acc))
//  [python]
def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)

        loss.backward()
        #Apply Nearest Orthogonal Gradient Mapping
        if args.nog:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad
            else:
                weight_grad = net.module.conv1.weight.grad
            B,C,H,W = weight_grad.size()
            u,s,vh=torch.linalg.svd(weight_grad.view(B,C*H*W),full_matrices=False)
            s[:] = 1
            if args.ow:
                net.module.conv1.random_conv_filter.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
            else:
                net.module.conv1.weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
        #Apply Optimal Learnig Rate
        if args.olr:
            if args.ow:
                weight_grad = net.module.conv1.random_conv_filter.grad.view(-1, 1)
                weight = net.module.conv1.random_conv_filter.view(-1,1)
            else:
                weight_grad = net.module.conv1.weight.grad.view(-1, 1)
                weight = net.module.conv1.weight.view(-1, 1)
            wtw = weight.t().mm(weight)
            ltw=weight_grad.t().mm(weight)
            ltl=weight_grad.t().mm(weight_grad)
            lr = (wtw*ltw)/(wtw*ltl+2*ltw*ltw)
            if lr>0 and lr < optimizer.param_groups[1]['lr']:
                optimizer.param_groups[0]['lr'] = lr
            else:
                optimizer.param_groups[0]['lr'] = optimizer.param_groups[1]['lr']
        optimizer.step()


        if batch_idx % 500 == 0:
            writer.add_scalar('loss/train_loss', loss.item(), epoch * len(trainloader) + batch_idx + 1)
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    acc = 100. * correct / total
    writer.add_scalar('loss/train_loss_avg', train_loss / len(trainloader), epoch)
    writer.add_scalar('train/accuracy', acc, epoch)
    writer.add_scalar('train/error', 100 - acc, epoch)
    print('current training error {}'.format(100 - acc))
//  [python]
def low_rank_process(Wo, hatWr, Lhr, args):
    # invLhr = torch.linalg.inv(Lhr)
    # assert torch.isfinite(invLhr).all()

    U, S, Vh = torch.svd((Wo - hatWr) @ Lhr)
    A = U[:, :args.lora_rank]
    # B = torch.diag(svdRZ.S[:args.lora_rank]) @ svdRZ.Vh[:args.lora_rank] @ invLhr
    B = torch.linalg.solve_triangular(
        Lhr,
        torch.diag(S[:args.lora_rank]) @ Vh[:args.lora_rank],
        upper=False,
        left=False)
    assert torch.isfinite(A).all() and torch.isfinite(B).all()

    U, S, Vh = torch.svd(B)
    A = (A @ U @ torch.diag(S.sqrt())).half()
    B = (torch.diag(S.sqrt()) @ Vh).half()

    hatWr = hatWr.to(A.device) + \
        (A @ B).to(torch.float64 if args.use_fp64 else torch.float32)
    return hatWr, A, B
//  [python]
def low_rank_process(Wo, hatWr, Lhr, args):
    # invLhr = torch.linalg.inv(Lhr)
    # assert torch.isfinite(invLhr).all()

    svdRZ = torch.linalg.svd((Wo - hatWr) @ Lhr, full_matrices=False)
    A = svdRZ.U[:, :args.lora_rank]
    # B = torch.diag(svdRZ.S[:args.lora_rank]) @ svdRZ.Vh[:args.lora_rank] @ invLhr
    B = torch.linalg.solve_triangular(
        Lhr,
        torch.diag(svdRZ.S[:args.lora_rank]) @ svdRZ.Vh[:args.lora_rank],
        upper=False,
        left=False)
    assert torch.isfinite(A).all() and torch.isfinite(B).all()

    svdB = torch.linalg.svd(B, full_matrices=False)
    A = (A @ svdB.U @ torch.diag(svdB.S.sqrt())).half()
    B = (torch.diag(svdB.S.sqrt()) @ svdB.Vh).half()

    hatWr = hatWr.to(A.device) + \
        (A @ B).to(torch.float64 if args.use_fp64 else torch.float32)
    return hatWr, A, B
//  [python]
def low_rank_preprocess(Wr, Hr, Lhr, args):
    dtype_ = torch.float64 if args.use_fp64 else torch.float32
    if args.full_svd:
        U, S, Vh = torch.svd(Wr.to(torch.float64) @ Lhr.to(torch.float64))
        svdZ = torch.linalg.SVD(Wr.to(torch.float64) @ Lhr.to(torch.float64), full_matrices=False)
        Hr -= (Lhr.to(torch.float64) @ Vh.T[:, :args.lora_rank] @ \
                   Vh[:args.lora_rank] @ Lhr.to(torch.float64).T).to(dtype_)
        Hr += torch.diag(Hr).mean() * args.sigma_reg2 * \
            torch.eye(Hr.shape[0], device=Hr.device, dtype=Hr.dtype)
        Wr -= (U[:, :args.lora_rank] @ U.T[:args.lora_rank] @ Wr.to(
            torch.float64)).to(dtype_)
    else:
        U_lrz, S_lrz, V_lrz = torch.svd_lowrank(
            Wr.to(torch.float64) @ Lhr.to(torch.float64),
            q=2 * args.lora_rank,
            niter=10)
        U_lrz = U_lrz[:, :args.lora_rank]
        V_lrz = V_lrz[:, :args.lora_rank]
        Hr -= (Lhr.to(torch.float64) @ V_lrz @ V_lrz.T @ Lhr.to(
            torch.float64).T).to(dtype_)
        Hr += torch.diag(Hr).mean() * args.sigma_reg2 * \
            torch.eye(Hr.shape[0], device=Hr.device, dtype=Hr.dtype)
        Wr -= (U_lrz @ U_lrz.T @ Wr.to(torch.float64)).to(dtype_)
    return Wr, Hr
//  [python]
def low_rank_preprocess(Wr, Hr, Lhr, args):
    dtype_ = torch.float64 if args.use_fp64 else torch.float32
    if args.full_svd:
        svdZ = torch.linalg.svd(Wr.to(torch.float64) @ Lhr.to(torch.float64),
                                full_matrices=False)
        Hr -= (Lhr.to(torch.float64) @ svdZ.Vh.T[:, :args.lora_rank] @ \
                   svdZ.Vh[:args.lora_rank] @ Lhr.to(torch.float64).T).to(dtype_)
        Hr += torch.diag(Hr).mean() * args.sigma_reg2 * \
            torch.eye(Hr.shape[0], device=Hr.device, dtype=Hr.dtype)
        Wr -= (svdZ.U[:, :args.lora_rank] @ svdZ.U.T[:args.lora_rank] @ Wr.to(
            torch.float64)).to(dtype_)
    else:
        U_lrz, S_lrz, V_lrz = torch.svd_lowrank(
            Wr.to(torch.float64) @ Lhr.to(torch.float64),
            q=2 * args.lora_rank,
            niter=10)
        U_lrz = U_lrz[:, :args.lora_rank]
        V_lrz = V_lrz[:, :args.lora_rank]
        Hr -= (Lhr.to(torch.float64) @ V_lrz @ V_lrz.T @ Lhr.to(
            torch.float64).T).to(dtype_)
        Hr += torch.diag(Hr).mean() * args.sigma_reg2 * \
            torch.eye(Hr.shape[0], device=Hr.device, dtype=Hr.dtype)
        Wr -= (U_lrz @ U_lrz.T @ Wr.to(torch.float64)).to(dtype_)
    return Wr, Hr
//  [python]
def _fit_full(self, x):
        """Fit the model by computing full SVD on X."""
        x = torch.tensor(x,dtype=torch.float32).to(self.device)
        n_samples, n_features = x.shape
        X_center = x-torch.mean(x, axis=0)   # Center data
        U, s, Vt = torch.svd(X_center)
        self.__Vt = Vt
        explained_variance = (s ** 2) / (n_samples - 1)  # Get variance explained by singular values
        total_var = explained_variance.sum()
        explained_variance_ratio = explained_variance / total_var  # caculate the compressed rate
        if self.n_component_ratio > 0 and self.n_component_ratio < 1:
            ratio_cumsum = stable_cumsum(explained_variance_ratio.cpu().numpy())
            self.n_components = np.searchsorted(ratio_cumsum, self.n_component_ratio, side="right") + 1
        elif type(self.n_component_ratio) == int and self.n_component_ratio < n_features:
            self.n_components = self.n_component_ratio
        else:
            AttributeError("n_components error!")
        self.explained_variance_ratio=explained_variance_ratio[:self.n_components]
        return U, s, Vt
//  [python]
def _fit_full(self, x):
        """Fit the model by computing full SVD on X."""
        x = torch.tensor(x,dtype=torch.float32).to(self.device)
        n_samples, n_features = x.shape
        X_center = x-torch.mean(x, axis=0)   # Center data
        U, s, Vt = torch.linalg.svd(X_center, full_matrices=False)
        self.__Vt = Vt
        explained_variance = (s ** 2) / (n_samples - 1)  # Get variance explained by singular values
        total_var = explained_variance.sum()
        explained_variance_ratio = explained_variance / total_var  # caculate the compressed rate
        if self.n_component_ratio > 0 and self.n_component_ratio < 1:
            ratio_cumsum = stable_cumsum(explained_variance_ratio.cpu().numpy())
            self.n_components = np.searchsorted(ratio_cumsum, self.n_component_ratio, side="right") + 1
        elif type(self.n_component_ratio) == int and self.n_component_ratio < n_features:
            self.n_components = self.n_component_ratio
        else:
            AttributeError("n_components error!")
        self.explained_variance_ratio=explained_variance_ratio[:self.n_components]
        return U, s, Vt
//  [python]
def batch_weighted_procrustes(X, Y, w,  eps=np.finfo(np.float32).eps):
  """
  X: torch tensor B x N x 3; NOCS
  Y: torch tensor B x N x 3; Input
  w: torch tensor B x N x 1
  """
  # https://ieeexplore.ieee.org/document/88573
  assert len(X) == len(Y)
  B = X.size(0)
  W1 = torch.abs(w).sum(dim=1, keepdim=True)
  w_norm = w / (W1 + eps)
  mux = (w_norm * X).sum(dim=1, keepdim=True)
  muy = (w_norm * Y).sum(dim=1, keepdim=True)  # torch geometric can get these quantities via mean-pool

  # Use CPU for small arrays
  Sxy = (Y.cpu() - muy.cpu()).transpose(1, 2) @ (w_norm.cpu() * (X.cpu() - mux.cpu()))
  Sxy = Sxy.double()
  Sxy += (torch.rand(3, 3) * eps).double()  # for numerical stability
  # mean-centering is easy via indexing. batchwise-matrix-mul

  U, S, V = torch.svd(Sxy)
  S = torch.eye(3).unsqueeze(0).repeat(B, 1, 1).to(Sxy)
  inv_det = torch.where((U.det() * V.det()) < 0, -torch.ones(B).to(U), torch.ones(B).to(Sxy))
  S[:, -1, -1] = inv_det

  R = U @ S @ V.transpose(1, 2)
  R = R.float().to(X.device)
  t = (muy.squeeze() - (R @ mux.transpose(1, 2)).squeeze()).float().to(X.device)

  return R, t
//  [python]
def batch_weighted_procrustes(X, Y, w,  eps=np.finfo(np.float32).eps):
  """
  X: torch tensor B x N x 3; NOCS
  Y: torch tensor B x N x 3; Input
  w: torch tensor B x N x 1
  """
  # https://ieeexplore.ieee.org/document/88573
  assert len(X) == len(Y)
  B = X.size(0)
  W1 = torch.abs(w).sum(dim=1, keepdim=True)
  w_norm = w / (W1 + eps)
  mux = (w_norm * X).sum(dim=1, keepdim=True)
  muy = (w_norm * Y).sum(dim=1, keepdim=True)  # torch geometric can get these quantities via mean-pool

  # Use CPU for small arrays
  Sxy = (Y.cpu() - muy.cpu()).transpose(1, 2) @ (w_norm.cpu() * (X.cpu() - mux.cpu()))
  Sxy = Sxy.double()
  Sxy += (torch.rand(3, 3) * eps).double()  # for numerical stability
  # mean-centering is easy via indexing. batchwise-matrix-mul

  U, D, V = torch.linalg.svd(Sxy)
  S = torch.eye(3).unsqueeze(0).repeat(B, 1, 1).to(Sxy)
  inv_det = torch.where((U.det() * V.det()) < 0, -torch.ones(B).to(U), torch.ones(B).to(Sxy))
  S[:, -1, -1] = inv_det

  R = U @ S @ V.transpose(1, 2)
  R = R.float().to(X.device)
  t = (muy.squeeze() - (R @ mux.transpose(1, 2)).squeeze()).float().to(X.device)

  return R, t
//  [python]
def batched_sym_matrix_pow(matrices: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrices: A batch of matrices.
        p: Power.
        positive_definite: If positive definite
    Returns:
        Power of each matrix in the batch.
    """
    # vals, vecs = torch.linalg.eigh(matrices)
    # SVD is much faster than  vals, vecs = torch.linalg.eigh(matrices) for large batches.
    vecs, vals, _ = torch.svd(matrices)
    good = vals > vals.max(-1, True).values * vals.size(-1) * torch.finfo(vals.dtype).eps
    vals = vals.pow(p).where(good, torch.zeros((), device=matrices.device, dtype=matrices.dtype))
    matrix_power = (vecs * vals.unsqueeze(-2)) @ torch.transpose(vecs, -2, -1)
    return matrix_power
//  [python]
def batched_sym_matrix_pow(matrices: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrices: A batch of matrices.
        p: Power.
        positive_definite: If positive definite
    Returns:
        Power of each matrix in the batch.
    """
    # vals, vecs = torch.linalg.eigh(matrices)
    # SVD is much faster than  vals, vecs = torch.linalg.eigh(matrices) for large batches.
    vecs, vals, _ = torch.linalg.svd(matrices)
    good = vals > vals.max(-1, True).values * vals.size(-1) * torch.finfo(vals.dtype).eps
    vals = vals.pow(p).where(good, torch.zeros((), device=matrices.device, dtype=matrices.dtype))
    matrix_power = (vecs * vals.unsqueeze(-2)) @ torch.transpose(vecs, -2, -1)
    return matrix_power
//  [python]
def to_mps(self,):
        mps = [t.array.reshape(1, 2, 1) for t in self.circuits_tn[:self.n]]
        for tensor in self.circuits_tn[self.n:]:
            array, inds = tensor.array, tensor.inds
            if len(inds) == 2:
                q = int(inds[0].split('-')[1])
                mps[q] = torch.einsum('abc,bd->adc', mps[q], array)
            elif len(inds) == 4:
                q1, q2 = int(inds[0].split('-')[1]), int(inds[1].split('-')[1])
                u, s, vh = torch.svd(array.permute(0, 2, 1, 3).reshape(4, 4))
                mq1 = (u @ torch.diag(torch.sqrt(s).to(self.dtype))).reshape(2, 2, -1)
                mq2 = (torch.diag(torch.sqrt(s).to(self.dtype)) @ vh).reshape(-1, 2, 2)
                mps[q1] = torch.einsum('abc, bde->adec', mps[q1], mq1).reshape(
                    mps[q1].shape[0], 2, -1
                )
                mps[q2] = torch.einsum('abc, ebd->eadc', mps[q2], mq2).reshape(
                    -1, 2, mps[q2].shape[2]
                )
            
        return mps
//  [python]
def to_mps(self,):
        mps = [t.array.reshape(1, 2, 1) for t in self.circuits_tn[:self.n]]
        for tensor in self.circuits_tn[self.n:]:
            array, inds = tensor.array, tensor.inds
            if len(inds) == 2:
                q = int(inds[0].split('-')[1])
                mps[q] = torch.einsum('abc,bd->adc', mps[q], array)
            elif len(inds) == 4:
                q1, q2 = int(inds[0].split('-')[1]), int(inds[1].split('-')[1])
                u, s, vh = torch.linalg.svd(array.permute(0, 2, 1, 3).reshape(4, 4))
                mq1 = (u @ torch.diag(torch.sqrt(s).to(self.dtype))).reshape(2, 2, -1)
                mq2 = (torch.diag(torch.sqrt(s).to(self.dtype)) @ vh).reshape(-1, 2, 2)
                mps[q1] = torch.einsum('abc, bde->adec', mps[q1], mq1).reshape(
                    mps[q1].shape[0], 2, -1
                )
                mps[q2] = torch.einsum('abc, ebd->eadc', mps[q2], mq2).reshape(
                    -1, 2, mps[q2].shape[2]
                )
            
        return mps
//  [python]
def custom_pca_image(img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    This will do a standard SVD on a 2D projections of `img`. `U @ S`
    is reshaped into a 3D tensor to match the shape im img at the start.
    MUST BE IN (ROWS, COLS, CHANNELS)

    Parameters
    ----------
    img: torch.Tensor
        MUST BE IN (ROWS, COLS, CHANNELS)

    Returns
    -------
    v_pca: torch.Tensor
        the vh from torch.linalg.svd(img, (nr * nc, p), full_matrices=False)
    pc: torch.Tensor
        U @ s reshaped into the number of rows and columns of the input img
    """
    nr, nc, p = img.shape
    # y_w -> h x w X c
    im1 = torch.reshape(img, (nr * nc, p))
    u, s, v_pca = torch.svd(im1)
    # need to modify u and s
    pc = torch.matmul(u, torch.diag(s))
    pc = pc.reshape((nc, nr, p))
    return v_pca, pc
//  [python]
def custom_pca_image(img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    This will do a standard SVD on a 2D projections of `img`. `U @ S`
    is reshaped into a 3D tensor to match the shape im img at the start.
    MUST BE IN (ROWS, COLS, CHANNELS)

    Parameters
    ----------
    img: torch.Tensor
        MUST BE IN (ROWS, COLS, CHANNELS)

    Returns
    -------
    v_pca: torch.Tensor
        the vh from torch.linalg.svd(img, (nr * nc, p), full_matrices=False)
    pc: torch.Tensor
        U @ s reshaped into the number of rows and columns of the input img
    """
    nr, nc, p = img.shape
    # y_w -> h x w X c
    im1 = torch.reshape(img, (nr * nc, p))
    u, s, v_pca = torch.linalg.svd(im1, full_matrices=False)
    # need to modify u and s
    pc = torch.matmul(u, torch.diag(s))
    pc = pc.reshape((nc, nr, p))
    return v_pca, pc
//  [python]
def build_usv_from_weight(self, weight: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        ### differentiable feature is gauranteed
        U, S, V = weight.data.svd(some=False)
        V = V.transpose(-2, -1).contiguous()
        self.U.data.copy_(U)
        self.S.data.copy_(S)
        self.V.data.copy_(V)
        return U, S, V
//  [python]
def build_usv_from_weight(self, weight: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        ### differentiable feature is gauranteed
        # U, S, V = weight.data.svd(some=False) # old version is not precise
        # V = V.transpose(-2, -1).contiguous()
        U, S, V = torch.linalg.svd(weight.data, full_matrices=True, driver="gesvd")  # must use QR decomposition
        self.U.data.copy_(U)
        self.S.data.copy_(S)
        self.V.data.copy_(V)
        return U, S, V
//  [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 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.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.solve(constraint_mat[:, :, :10], b)[0]

        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.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