jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_svd_lowrank(self, device):
        import torch
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        dtype = torch.double

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.transpose(-2, -1))
            self.assertEqual(A, a)

            # check if svd_lowrank produces same singular values as torch.linalg.svd
            U, S, V = torch.linalg.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.transpose(-2, -1).matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.transpose(-2, -1).matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)
//  [python]
def mx_svd(mx, topk):
    U, D, V = torch.svd(mx) 
    #topk x m
    evecs = U.t()[-topk-1:-1]
    row1 = evecs[-1, :]
    idx = torch.argsort(row1)
    evecs = evecs[:, idx]
    return evecs
//  [python]
def mx_svd(mx, topk):
    U, S, V = torch.linalg.svd(mx)
    evecs = U.t()[-topk-1:-1]
    row1 = evecs[-1, :]
    idx = torch.argsort(row1)
    evecs = evecs[:, idx]
    return evecs
//  [python]
def soft_dimension(features):
    """Continuous approximation of the rank of a (N, D) sample.

    Let "s" denote the (D,) vector of eigenvalues of Cov,
    the (D, D) covariance matrix of the sample "features".
    Then,
        R(features) = \sum_i sqrt(s_i) / \max_i sqrt(s_i)

    This quantity encodes the number of PCA components that would be
    required to describe the sample with a good precision.
    It is equal to D if the sample is isotropic, but is generally much lower.

    Up to the re-normalization by the largest eigenvalue,
    this continuous pseudo-rank is equal to the nuclear norm of the sample.
    """

    nfeat = features.shape[-1]
    features = features.view(-1, nfeat)
    x = features - torch.mean(features, dim=0, keepdim=True)
    cov = x.T @ x
    try:
        u, s, v = torch.svd(cov)
        R = s.sqrt().sum() / s.sqrt().max()
    except:
        return -1
    return R.item()
//  [python]
def soft_dimension(features):
    """Continuous approximation of the rank of a (N, D) sample.

    Let "s" denote the (D,) vector of eigenvalues of Cov,
    the (D, D) covariance matrix of the sample "features".
    Then,
        R(features) = \sum_i sqrt(s_i) / \max_i sqrt(s_i)

    This quantity encodes the number of PCA components that would be
    required to describe the sample with a good precision.
    It is equal to D if the sample is isotropic, but is generally much lower.

    Up to the re-normalization by the largest eigenvalue,
    this continuous pseudo-rank is equal to the nuclear norm of the sample.
    """

    nfeat = features.shape[-1]
    features = features.view(-1, nfeat)
    x = features - torch.mean(features, dim=0, keepdim=True)
    cov = x.T @ x
    try:
        u, s, v = torch.linalg.svd(cov)
        R = s.sqrt().sum() / s.sqrt().max()
    except:
        return -1
    return R.item()
//  [python]
def mx_inv(mx):
    if isinstance(mx, torch_sparse.tensor.SparseTensor):
        mx = mx.to_dense()
    U, D, V = torch.svd(mx)
    eps = 0.009
    D_min = torch.min(D)
    if D_min < eps:
        D_1 = torch.zeros_like(D)
        D_1[D>D_min] = 1/D[D>D_min]
    else:
        D_1 = 1/D
    #D_1 = 1 / D #.clamp(min=0.005)
    
    return U @ D_1.diag() @ V.t()
//  [python]
def mx_inv(mx):
    if isinstance(mx, torch_sparse.tensor.SparseTensor):
        mx = mx.to_dense()
    U, S, V = torch.linalg.svd(mx)
    eps = 0.009
    D_min = torch.min(S)
    if D_min < eps:
        D_1 = torch.zeros_like(S)
        D_1[S>D_min] = 1/S[S>D_min]
    else:
        D_1 = 1/S
    #D_1 = 1 / D #.clamp(min=0.005)
    
    return U @ D_1.diag() @ V.t()
//  [python]
def mx_inv_sqrt(mx):
    # singular values need to be distinct for backprop
    U, D, V = torch.svd(mx)
    D_min = torch.min(D)
    eps = 0.009
    if D_min < eps:
        D_1 = torch.zeros_like(D)        
        D_1[D>D_min] = 1/D[D>D_min] #.sqrt()
    else:
        D_1 = 1/D #.sqrt()
    #D_1 = 1 / D.clamp(min=0.005).sqrt()
    return U @ D_1.sqrt().diag() @ V.t(), U @ D_1.diag() @ V.t()
//  [python]
def mx_inv_sqrt(mx):
    # singular values need to be distinct for backprop
    U, S, V = torch.linalg.svd(mx)
    D_min = torch.min(S)
    eps = 0.009
    if D_min < eps:
        D_1 = torch.zeros_like(S)        
        D_1[S>D_min] = 1/S[S>D_min] #.sqrt()
    else:
        D_1 = 1/S #.sqrt()
    #D_1 = 1 / S.clamp(min=0.005).sqrt()
    return U @ D_1.sqrt().diag() @ V.t(), U @ D_1.diag() @ V.t()
//  [python]
def bench_speed(N, H, W):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    torch.manual_seed(0)
    a = torch.randn(N, H, W).cuda()
    b = a.clone().cuda()
    torch.cuda.synchronize()

    start.record()
    for i in range(100):
        U, S, V = svd(a)
    end.record()
    torch.cuda.synchronize()
    t = start.elapsed_time(end) / 100
    print("Perform batched SVD on a {}x{}x{} matrix: {} ms".format(N, H, W, t))

    start.record()
    U, S, V = torch.svd(b, some=True, compute_uv=True)
    end.record()
    torch.cuda.synchronize()
    t = start.elapsed_time(end)
    print("Perform torch.svd on a {}x{}x{} matrix: {} ms".format(N, H, W, t))
//  [python]
def bench_speed(N, H, W):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    torch.manual_seed(0)
    a = torch.randn(N, H, W).cuda()
    b = a.clone().cuda()
    torch.cuda.synchronize()

    start.record()
    for i in range(100):
        U, S, V = torch.linalg.svd(a)
    end.record()
    torch.cuda.synchronize()
    t = start.elapsed_time(end) / 100
    print("Perform batched SVD on a {}x{}x{} matrix: {} ms".format(N, H, W, t))

    start.record()
    U, S, V = torch.linalg.svd(b)
    end.record()
    torch.cuda.synchronize()
    t = start.elapsed_time(end)
    print("Perform torch.svd on a {}x{}x{} matrix: {} ms".format(N, H, W, t))
//  [python]
def rigid_transform_3d(p1, p2):
    center_p1 = torch.mean(p1, dim=0, keepdim=True)
    center_p2 = torch.mean(p2, dim=0, keepdim=True)

    pp1 = p1 - center_p1
    pp2 = p2 - center_p2

    h = torch.mm(pp1.t(), pp2)
    u, _, v = torch.svd(h)
    r = torch.mm(v, u.T)

    # reflection
    if torch.linalg.det(r) < 0:
        v[2, :] *= -1
        r = torch.mm(v.t(), u.t())

    t = torch.mm(-r, center_p1.t()) + center_p2.t()

    return r, t
//  [python]
def rigid_transform_3d(p1, p2):
    center_p1 = torch.mean(p1, dim=0, keepdim=True)
    center_p2 = torch.mean(p2, dim=0, keepdim=True)

    pp1 = p1 - center_p1
    pp2 = p2 - center_p2

    h = torch.mm(pp1.t(), pp2)
    u, s, v = torch.linalg.svd(h)
    r = torch.mm(v, u.T)

    # reflection
    if torch.linalg.det(r) < 0:
        v[2, :] *= -1
        r = torch.mm(v.t(), u.t())

    t = torch.mm(-r, center_p1.t()) + center_p2.t()

    return r, t
//  [python]
def svd(M):
    """
    Singular Value Decomposition wrapper, using efficient batch decomposition on GPU when available.

    Args:
        M (BxMxN tensor): batch of real matrices.
    Returns:
        (U,D,V) decomposition, such as :math:`M = U @ diag(D) @ V^T`.
    """
    if M.is_cuda and M.shape[1] < 32 and M.shape[2] < 32:
        return _fast_gpu_svd(M)
    else:
        return torch.svd(M)
//  [python]
def svd(M):
    """
    Singular Value Decomposition wrapper, using efficient batch decomposition on GPU when available.

    Args:
        M (BxMxN tensor): batch of real matrices.
    Returns:
        (U,D,V) decomposition, such as :math:`M = U @ diag(D) @ V^T`.
    """
    if M.is_cuda and M.shape[1] < 32 and M.shape[2] < 32:
        return _fast_gpu_svd(M)
    else:
        return torch.linalg.svd(M)
//  [python]
def nearestPD(A):
    """ dont use this it doesn't work"""
    B = A + A.T / 2
    u, s, v = torch.svd(B)
    H = torch.matmul(v.T, torch.matmul(s.diag_embed(), v))
    A2 = (B + H) / 2
    A3 = (A2 + A2.t()) / 2
    if isPD(A3):
        return A3

    spacing = torch.finfo(A3.dtype).eps
    I = torch.eye(A.size(0), device=A.device)
    k = 1
    while not isPD(A3):
        val, vecs = A3.eig()
        mineig = val[0].min()
        A3 += I * (-mineig * k ** 2 + spacing)
        k += 1

    return A3
//  [python]
def nearestPD(A):
    """ dont use this it doesn't work"""
    B = A + A.T / 2
    u, s, v = torch.linalg.svd(B)
    H = torch.matmul(v.T, torch.matmul(s.diag_embed(), v))
    A2 = (B + H) / 2
    A3 = (A2 + A2.t()) / 2
    if isPD(A3):
        return A3

    spacing = torch.finfo(A3.dtype).eps
    I = torch.eye(A.size(0), device=A.device)
    k = 1
    while not isPD(A3):
        val, vecs = A3.eig()
        mineig = val[0].min()
        A3 += I * (-mineig * k ** 2 + spacing)
        k += 1

    return A3
//  [python]
def pca_lmark_lrw():
    k = 20
    lmark_root_path = '/mnt/ssd0/dat/lchen63/lrw/data/landmark1d/'
    _file = open( '/home/cxu-serve/p1/lchen63/lrw/pickle/region_16_warp_gt_test2.pkl', "rb")
    train_data = pickle.load(_file)
    shuffle(train_data)
    print (len(train_data))
    landmarks = []
    for index in range(len(train_data)):
        if index == 80000:
            
            break

        lmark_path = os.path.join(lmark_root_path ,train_data[index][0][:-8] + '.npy') 
        t_lmark = np.load(lmark_path)
        landmarks.append(torch.from_numpy(t_lmark[random.randint(0, 28)]))
        print (lmark_path)
    landmarks = torch.stack(landmarks,dim= 0)
    print (landmarks.shape)
    landmarks = landmarks.view(landmarks.size(0), 136)
    mean = torch.mean(landmarks,0)
    print (mean)

    np.save('../basics/mean_lrw_region.npy', mean.numpy())
    landmarks = landmarks - mean.expand_as(landmarks)

    U,S,V  = torch.svd(torch.t(landmarks))
    np.save('../basics/U_lrw_region.npy', U.numpy())
//  [python]
def pca_lmark_lrw():
    k = 20
    lmark_root_path = '/mnt/ssd0/dat/lchen63/lrw/data/landmark1d/'
    _file = open( '/home/cxu-serve/p1/lchen63/lrw/pickle/region_16_warp_gt_test2.pkl', "rb")
    train_data = pickle.load(_file)
    shuffle(train_data)
    print (len(train_data))
    landmarks = []
    for index in range(len(train_data)):
        if index == 80000:
            
            break

        lmark_path = os.path.join(lmark_root_path ,train_data[index][0][:-8] + '.npy') 
        t_lmark = np.load(lmark_path)
        landmarks.append(torch.from_numpy(t_lmark[random.randint(0, 28)]))
        print (lmark_path)
    landmarks = torch.stack(landmarks,dim= 0)
    print (landmarks.shape)
    landmarks = landmarks.view(landmarks.size(0), 136)
    mean = torch.mean(landmarks,0)
    print (mean)

    np.save('../basics/mean_lrw_region.npy', mean.numpy())
    landmarks = landmarks - mean.expand_as(landmarks)

    U,S,V  = torch.linalg.svd(torch.t(landmarks))
    np.save('../basics/U_lrw_region.npy', U.numpy())
//  [python]
def PCA_svd(X, k: int):
    m, t, p = X.size()
    X = X.reshape(m * t, p)
    n, p = X.shape
    X_center = X - X.mean(1).view(n, 1)
    u, s, v = torch.svd(X_center)
    components = v[:k].T
    X = X @ components
    return X.reshape(m, t, k)
//  [python]
def PCA_svd(X, k: int):
    m, t, p = X.size()
    X = X.reshape(m * t, p)
    n, p = X.shape
    X_center = X - X.mean(1).view(n, 1)
    u, s, v = torch.linalg.svd(X_center)
    components = v[:k].T
    X = X @ components
    return X.reshape(m, t, k)
//  [python]
def pca(Y, embedding_dim):
    """PCA embedding of a data matrix.

    Arguments
    ---------
    Y: torch.Tensor(shape=(n, k))
        data matrix, with n >= k
    embedding_dim: int
        the number of eigenvectors to retrieve, the embedding dimension;
        must be <= k

    Returns
    -------
    torch.Tensor(shape=(n, embedding_dim))
        The top embedding_dim eigenvectors of YY^T, scaled by sqrt(n)
    """
    n = Y.shape[0]
    embedding_dim = embedding_dim.int()
    min_dim = min([n, Y.shape[1]])
    if embedding_dim > min_dim:
        raise ValueError(
            "Embedding dimension must be at most minimum dimension of Y"
        )

    # PCA requires the data to be centered.
    Y = Y - Y.mean(axis=0)[None, :]
    U, _, _ = torch.svd(Y)
    return np.sqrt(float(n)) * U[:, :embedding_dim]
//  [python]
def pca(Y, embedding_dim):
    """PCA embedding of a data matrix.

    Arguments
    ---------
    Y: torch.Tensor(shape=(n, k))
        data matrix, with n >= k
    embedding_dim: int
        the number of eigenvectors to retrieve, the embedding dimension;
        must be <= k

    Returns
    -------
    torch.Tensor(shape=(n, embedding_dim))
        The top embedding_dim eigenvectors of YY^T, scaled by sqrt(n)
    """
    n = Y.shape[0]
    embedding_dim = embedding_dim.int()
    min_dim = min([n, Y.shape[1]])
    if embedding_dim > min_dim:
        raise ValueError(
            "Embedding dimension must be at most minimum dimension of Y"
        )

    # PCA requires the data to be centered.
    Y = Y - Y.mean(axis=0)[None, :]
    U, _, _ = torch.linalg.svd(Y)
    return np.sqrt(float(n)) * U[:, :embedding_dim]
//  [python]
def SVDreduction(ndatas,K):
    # ndatas = torch.linear.qr(datas.permute(0, 2, 1),'reduced').R
    # ndatas = ndatas.permute(0, 2, 1)
    _,s,v = torch.svd(ndatas)
    ndatas = ndatas.matmul(v[:,:,:K])

    return ndatas
//  [python]
def SVDreduction(ndatas,K):
    U,S,V = torch.linalg.svd(ndatas)
    ndatas = ndatas.matmul(V[:,:,:K])

    return ndatas
//  [python]
def cca(data_tensor_x, data_tensor_y, batch_size=100):
    '''Perform canonical correlation analysis for two data tensors.

    Arguments:
        data_tensor_x (Tensor): contains the first data tensor
        data_tensor_y (Tensor): contains the second data tensor
        batch_size (int): specify a batch size for the CCA calculation
    '''
    loader = _DataLoader(
        _TensorDataset(data_tensor_x, data_tensor_y),
        batch_size=batch_size)
    x_mean, y_mean = get_mean(loader)
    cxx, cxy, cyy = get_covariance(loader, x_mean, y_mean)
    ixx = get_sqrt_inverse(cxx)
    iyy = get_sqrt_inverse(cyy)
    return _torch.svd(_torch.mm(_torch.mm(ixx, cxy), iyy))
//  [python]
def cca(data_tensor_x, data_tensor_y, batch_size=100):
    '''Perform canonical correlation analysis for two data tensors.

    Arguments:
        data_tensor_x (Tensor): contains the first data tensor
        data_tensor_y (Tensor): contains the second data tensor
        batch_size (int): specify a batch size for the CCA calculation
    '''
    loader = _DataLoader(
        _TensorDataset(data_tensor_x, data_tensor_y),
        batch_size=batch_size)
    x_mean, y_mean = get_mean(loader)
    cxx, cxy, cyy = get_covariance(loader, x_mean, y_mean)
    ixx = get_sqrt_inverse(cxx)
    iyy = get_sqrt_inverse(cyy)
    return _torch.linalg.svd(_torch.mm(_torch.mm(ixx, cxy), iyy))
//  [python]
def benchmark(method):

    start_time = time.time()
    times = []

    for i in range(1):
        if method == 'svd':
            _result = torch.svd(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'inv':
            _result = torch.inverse(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'pinv':
            _result = u.pinv(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'pinverse':
            _result = torch.pinverse(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'eig':
            _result = torch.symeig(H, eigenvectors=True)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'svd':
            _result = torch.svd(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'solve':
            _result = torch.solve(S, H)
            open('/dev/null', 'w').write(str(_result[0]))
        else:
            assert False
        new_time = time.time()
        elapsed_time = 1000 * (new_time - start_time)
        print(f"{elapsed_time:8.2f}   {method}")
        start_time = new_time
        times.append(elapsed_time)
//  [python]
def benchmark(method):

    start_time = time.time()
    times = []

    for i in range(1):
        if method == 'svd':
            _result = torch.linalg.svd(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'inv':
            _result = torch.inverse(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'pinv':
            _result = u.pinv(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'pinverse':
            _result = torch.pinverse(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'eig':
            _result = torch.symeig(H, eigenvectors=True)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'svd':
            _result = torch.linalg.svd(H)
            open('/dev/null', 'w').write(str(_result[0]))
        elif method == 'solve':
            _result = torch.solve(S, H)
            open('/dev/null', 'w').write(str(_result[0]))
        else:
            assert False
        new_time = time.time()
        elapsed_time = 1000 * (new_time - start_time)
        print(f"{elapsed_time:8.2f}   {method}")
        start_time = new_time
        times.append(elapsed_time)
//  [python]
def __init__(self,encode_layers,bilinear_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        ELayers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            ELayers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                ELayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(ELayers)
        BLayers = OrderedDict()
        for layer_i in range(len(bilinear_layers)-1):
            BLayers["linear_{}".format(layer_i)] = nn.Linear(bilinear_layers[layer_i],bilinear_layers[layer_i+1])
            if layer_i != len(bilinear_layers)-2:
                BLayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.bilinear_net = nn.Sequential(BLayers)           
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, _, V = torch.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(bilinear_layers[-1],Nkoopman,bias=False)
//  [python]
def __init__(self,encode_layers,bilinear_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        ELayers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            ELayers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                ELayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(ELayers)
        BLayers = OrderedDict()
        for layer_i in range(len(bilinear_layers)-1):
            BLayers["linear_{}".format(layer_i)] = nn.Linear(bilinear_layers[layer_i],bilinear_layers[layer_i+1])
            if layer_i != len(bilinear_layers)-2:
                BLayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.bilinear_net = nn.Sequential(BLayers)           
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, S, V = torch.linalg.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(bilinear_layers[-1],Nkoopman,bias=False)
//  [python]
def rigid_transform_torch(A, B):

    assert A.size() == B.size()

    batch_size, num_rows, num_cols = A.size()
   
    # find mean column wise
    centroid_A = torch.mean(A.transpose(2,1).contiguous(), axis=1)
    centroid_B = torch.mean(B.transpose(2,1).contiguous(), axis=1)

    # ensure centroids are 3x1
    centroid_A = centroid_A.reshape(batch_size,num_rows,1)
    centroid_B = centroid_B.reshape(batch_size,num_rows,1)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = torch.matmul(Am, Bm.transpose(2,1).contiguous())

    # find rotation
    U, S, V = torch.svd(H)
    Z = torch.matmul(V,U.transpose(2,1).contiguous())
    # special reflection case
    d= (torch.linalg.det(Z) < 0).type(torch.int8)
    # -1/1 
    d=d*2-1
    Vc = V.clone()
    Vc[:,2,:]*=-d.view(batch_size,1)
    R = torch.matmul(Vc,U.transpose(2,1).contiguous())
   
    t = torch.matmul(-R, centroid_A)+centroid_B
    
    Trans=torch.cat((torch.cat((R,t),axis=2), torch.tensor([0,0,0,1]).repeat(batch_size,1).cuda().view(batch_size,1,4)),axis=1)

    return Trans
//  [python]
def rigid_transform_torch(A, B):

    assert A.size() == B.size()

    batch_size, num_rows, num_cols = A.size()
   
    # find mean column wise
    centroid_A = torch.mean(A.transpose(2,1).contiguous(), axis=1)
    centroid_B = torch.mean(B.transpose(2,1).contiguous(), axis=1)

    # ensure centroids are 3x1
    centroid_A = centroid_A.reshape(batch_size,num_rows,1)
    centroid_B = centroid_B.reshape(batch_size,num_rows,1)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = torch.matmul(Am, Bm.transpose(2,1).contiguous())

    # find rotation
    U, S, V = torch.linalg.svd(H)
    Z = torch.matmul(V,U.transpose(2,1).contiguous())
    # special reflection case
    d= (torch.linalg.det(Z) < 0).type(torch.int8)
    # -1/1 
    d=d*2-1
    Vc = V.clone()
    Vc[:,2,:]*=-d.view(batch_size,1)
    R = torch.matmul(Vc,U.transpose(2,1).contiguous())
   
    t = torch.matmul(-R, centroid_A)+centroid_B
    
    Trans=torch.cat((torch.cat((R,t),axis=2), torch.tensor([0,0,0,1]).repeat(batch_size,1).cuda().view(batch_size,1,4)),axis=1)

    return Trans
//  [python]
def __wct_core(self, cont_feat, styl_feat, noise):
        cFSize = cont_feat.size()
        c_mean = torch.mean(cont_feat, 1)  # c x (h x w)
        c_mean = c_mean.unsqueeze(1).expand_as(cont_feat)
        cont_feat = cont_feat - c_mean
        
        iden = torch.eye(cFSize[0])  # .double()
        if self.is_cuda:
            iden = iden.cuda()
        
        contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1]) + iden
        #contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1] -1) + iden
        # del iden
        c_u, c_e, c_v = torch.svd(contentConv, some=False)
        # c_e2, c_v = torch.eig(contentConv, True)
        # c_e = c_e2[:,0]
        
        k_c = cFSize[0]
        for i in range(cFSize[0] - 1, -1, -1):
            if c_e[i] >= 0.00001:
                k_c = i + 1
                break
        
        sFSize = styl_feat.size()
        s_mean = torch.mean(styl_feat, 1)
        styl_feat = styl_feat - s_mean.unsqueeze(1).expand_as(styl_feat)
        styleConv = torch.mm(styl_feat, styl_feat.t()).div(sFSize[1])
        #styleConv = torch.mm(styl_feat, styl_feat.t()).div(sFSize[1] - 1)
        s_u, s_e, s_v = torch.svd(styleConv, some=False)
        
        k_s = sFSize[0]
        for i in range(sFSize[0] - 1, -1, -1):
            if s_e[i] >= 0.00001:
                k_s = i + 1
                break
        
        c_d = (c_e[0:k_c]).pow(-0.5)
        step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d))
        step2 = torch.mm(step1, (c_v[:, 0:k_c].t()))
        whiten_cF = torch.mm(step2, cont_feat)

        s_d = (s_e[0:k_s]).pow(0.5)


        # noise svd
        if noise == True:
            noiseFeature1 = torch.randn(k_s,k_s).float()

            n_u, n_e, n_v = torch.svd(noiseFeature1, some=False)
            #n_v = noiseFeature1
            noiseFeature = torch.mm(torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), n_v.cuda()),(s_v[:, 0:k_s].t())), whiten_cF)
            originalFeature = torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), (s_v[:, 0:k_s].t())), whiten_cF)

            ratio_n = 1
            targetFeature = ratio_n * noiseFeature + (1 - ratio_n) * originalFeature

        else:
            targetFeature = torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), (s_v[:, 0:k_s].t())), whiten_cF)
        targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
        return targetFeature
//  [python]
def __wct_core(self, cont_feat, styl_feat, noise):
        cFSize = cont_feat.size()
        c_mean = torch.mean(cont_feat, 1)  # c x (h x w)
        c_mean = c_mean.unsqueeze(1).expand_as(cont_feat)
        cont_feat = cont_feat - c_mean
        
        iden = torch.eye(cFSize[0])  # .double()
        if self.is_cuda:
            iden = iden.cuda()
        
        contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1]) + iden
        #contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1] -1) + iden
        # del iden
        c_u, c_s, c_v = torch.linalg.svd(contentConv, full_matrices=False)
        # c_e2, c_v = torch.eig(contentConv, True)
        # c_e = c_e2[:,0]
        
        k_c = cFSize[0]
        for i in range(cFSize[0] - 1, -1, -1):
            if c_s[i] >= 0.00001:
                k_c = i + 1
                break
        
        sFSize = styl_feat.size()
        s_mean = torch.mean(styl_feat, 1)
        styl_feat = styl_feat - s_mean.unsqueeze(1).expand_as(styl_feat)
        styleConv = torch.mm(styl_feat, styl_feat.t()).div(sFSize[1])
        #styleConv = torch.mm(styl_feat, styl_feat.t()).div(sFSize[1] - 1)
        s_u, s_e, s_v = torch.linalg.svd(styleConv, full_matrices=False)
        
        k_s = sFSize[0]
        for i in range(sFSize[0] - 1, -1, -1):
            if s_e[i] >= 0.00001:
                k_s = i + 1
                break
        
        c_d = (c_s[0:k_c]).pow(-0.5)
        step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d))
        step2 = torch.mm(step1, (c_v[:, 0:k_c].t()))
        whiten_cF = torch.mm(step2, cont_feat)

        s_d = (s_e[0:k_s]).pow(0.5)


        # noise svd
        if noise == True:
            noiseFeature1 = torch.randn(k_s,k_s).float()

            n_u, n_e, n_v = torch.linalg.svd(noiseFeature1, full_matrices=False)
            #n_v = noiseFeature1
            noiseFeature = torch.mm(torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), n_v.cuda()),(s_v[:, 0:k_s].t())), whiten_cF)
            originalFeature = torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), (s_v[:, 0:k_s].t())), whiten_cF)

            ratio_n = 1
            targetFeature = ratio_n * noiseFeature + (1 - ratio_n) * originalFeature

        else:
            targetFeature = torch.mm(torch.mm(torch.mm(s_v[:, 0:k_s], torch.diag(s_d)), (s_v[:, 0:k_s].t())), whiten_cF)
        targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
        return targetFeature
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / batch_size  # d * d
        U, S, V = torch.svd(f_cov)
        diag = torch.diag(1.0 / torch.sqrt(S + self.eps))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / batch_size  # d * d
        U, S, V = torch.linalg.svd(f_cov)
        diag = torch.diag(1.0 / torch.sqrt(S + self.eps))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, _, V = torch.svd(S_non_zero)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, S, V = torch.linalg.svd(S_non_zero)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, _, V = torch.svd(S)

    rot_mat = torch.bmm(V, U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, S, V = torch.linalg.svd(S)

    rot_mat = torch.bmm(V, U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def get_render_error_in_q(args, model, sample_size, device, targets, rgbs, poses, batch_size=1):
    ''' use nerf render imgs instead of use real imgs '''
    model.eval()

    results = np.zeros((sample_size, 2))
    print("to be implement...")

    predict_pose_list = []
    gt_pose_list = []
    ang_error_list = []

    for i in range(sample_size):
        data = rgbs[i:i + 1].permute(0, 3, 1, 2)
        pose = poses[i:i + 1].reshape(batch_size, 12)

        data = data.to(device)  # input
        pose = pose.reshape((batch_size, 3, 4)).numpy()  # label

        # using SVD to make sure predict rotation is normalized rotation matrix
        with torch.no_grad():
            _, predict_pose = model(data)
            R_torch = predict_pose.reshape((batch_size, 3, 4))[:, :3, :3]  # debug
            predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()

            R = predict_pose[:, :3, :3]
            res = R @ np.linalg.inv(R)
            # print('R@np.linalg.inv(R):', res)

            u, s, v = torch.svd(R_torch)
            Rs = torch.matmul(u, v.transpose(-2, -1))
        predict_pose[:, :3, :3] = Rs[:, :3, :3].cpu().numpy()

        pose_q = transforms.matrix_to_quaternion(
            torch.Tensor(pose[:, :3, :3]))  # .cpu().numpy() # gnd truth in quaternion
        pose_x = pose[:, :3, 3]  # gnd truth position
        predicted_q = transforms.matrix_to_quaternion(
            torch.Tensor(predict_pose[:, :3, :3]))  # .cpu().numpy() # predict in quaternion
        predicted_x = predict_pose[:, :3, 3]  # predict position
        pose_q = pose_q.squeeze()
        pose_x = pose_x.squeeze()
        predicted_q = predicted_q.squeeze()
        predicted_x = predicted_x.squeeze()

        # Compute Individual Sample Error
        q1 = pose_q / torch.linalg.norm(pose_q)
        q2 = predicted_q / torch.linalg.norm(predicted_q)
        d = torch.abs(torch.sum(torch.matmul(q1, q2)))
        d = torch.clamp(d, -1., 1.)  # acos can only input [-1~1]
        theta = (2 * torch.acos(d) * 180 / math.pi).numpy()
        error_x = torch.linalg.norm(torch.Tensor(pose_x - predicted_x)).numpy()
        results[i, :] = [error_x, theta]
        # print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta))

        # save results for visualization
        predict_pose_list.append(predicted_x)
        gt_pose_list.append(pose_x)
        ang_error_list.append(theta)

    predict_pose_list = np.array(predict_pose_list)
    gt_pose_list = np.array(gt_pose_list)
    ang_error_list = np.array(ang_error_list)
    vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}

    median_result = np.median(results, axis=0)
    mean_result = np.mean(results, axis=0)

    # standard log
    print('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
    print('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))

    return
//  [python]
def get_render_error_in_q(args, model, sample_size, device, targets, rgbs, poses, batch_size=1):
    ''' use nerf render imgs instead of use real imgs '''
    model.eval()

    results = np.zeros((sample_size, 2))
    print("to be implement...")

    predict_pose_list = []
    gt_pose_list = []
    ang_error_list = []

    for i in range(sample_size):
        data = rgbs[i:i + 1].permute(0, 3, 1, 2)
        pose = poses[i:i + 1].reshape(batch_size, 12)

        data = data.to(device)  # input
        pose = pose.reshape((batch_size, 3, 4)).numpy()  # label

        # using SVD to make sure predict rotation is normalized rotation matrix
        with torch.no_grad():
            _, predict_pose = model(data)
            R_torch = predict_pose.reshape((batch_size, 3, 4))[:, :3, :3]  # debug
            predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()

            R = predict_pose[:, :3, :3]
            res = R @ np.linalg.inv(R)
            # print('R@np.linalg.inv(R):', res)

            u, s, v = torch.linalg.svd(R_torch)
            Rs = torch.matmul(u, v.transpose(-2, -1))
        predict_pose[:, :3, :3] = Rs[:, :3, :3].cpu().numpy()

        pose_q = transforms.matrix_to_quaternion(
            torch.Tensor(pose[:, :3, :3]))  # .cpu().numpy() # gnd truth in quaternion
        pose_x = pose[:, :3, 3]  # gnd truth position
        predicted_q = transforms.matrix_to_quaternion(
            torch.Tensor(predict_pose[:, :3, :3]))  # .cpu().numpy() # predict in quaternion
        predicted_x = predict_pose[:, :3, 3]  # predict position
        pose_q = pose_q.squeeze()
        pose_x = pose_x.squeeze()
        predicted_q = predicted_q.squeeze()
        predicted_x = predicted_x.squeeze()

        # Compute Individual Sample Error
        q1 = pose_q / torch.linalg.norm(pose_q)
        q2 = predicted_q / torch.linalg.norm(predicted_q)
        d = torch.abs(torch.sum(torch.matmul(q1, q2)))
        d = torch.clamp(d, -1., 1.)  # acos can only input [-1~1]
        theta = (2 * torch.acos(d) * 180 / math.pi).numpy()
        error_x = torch.linalg.norm(torch.Tensor(pose_x - predicted_x)).numpy()
        results[i, :] = [error_x, theta]
        # print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta))

        # save results for visualization
        predict_pose_list.append(predicted_x)
        gt_pose_list.append(pose_x)
        ang_error_list.append(theta)

    predict_pose_list = np.array(predict_pose_list)
    gt_pose_list = np.array(gt_pose_list)
    ang_error_list = np.array(ang_error_list)
    vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}

    median_result = np.median(results, axis=0)
    mean_result = np.mean(results, axis=0)

    # standard log
    print('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
    print('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))

    return
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """ 
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence 
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t 
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    # weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-6)

    # find mean of point cloud
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    # construct weight covariance matrix
    Weight = torch.diag_embed(weights)
    H = Am.permute(0, 2, 1) @ Weight @ Bm

    # find rotation
    U, S, Vt = torch.svd(H.cpu())
    U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1).to(A.device)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0,2,1) - R @ centroid_A.permute(0,2,1)
    # warp_A = transform(A, integrate_trans(R,t))
    # RMSE = torch.sum( (warp_A - B) ** 2, dim=-1).mean()
    return integrate_trans(R, t)
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """ 
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence 
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t 
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    # weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-6)

    # find mean of point cloud
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    # construct weight covariance matrix
    Weight = torch.diag_embed(weights)
    H = Am.permute(0, 2, 1) @ Weight @ Bm

    # find rotation
    U, S, Vt = torch.linalg.svd(H.cpu())
    U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1).to(A.device)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0,2,1) - R @ centroid_A.permute(0,2,1)
    # warp_A = transform(A, integrate_trans(R,t))
    # RMSE = torch.sum( (warp_A - B) ** 2, dim=-1).mean()
    return integrate_trans(R, t)
//  [python]
def _align_camera_extrinsics(
    cameras_src: "CamerasBase",
    cameras_tgt: "CamerasBase",
    estimate_scale: bool = True,
    eps: float = 1e-9,
):  # pragma: no cover
    """
    Get the global rotation R_A with svd of cov(RR^T):
        ```
        R_A R_i = R_i' for all i
        R_A [R_1 R_2 ... R_N] = [R_1' R_2' ... R_N']
        U, _, V = svd([R_1 R_2 ... R_N]^T [R_1' R_2' ... R_N'])
        R_A = (U V^T)^T
        ```
    """
    RRcov = torch.bmm(cameras_src.R, cameras_tgt.R.transpose(2, 1)).mean(0)
    U, _, V = torch.svd(RRcov)
    align_t_R = V @ U.t()

    """
    The translation + scale `T_A` and `s_A` is computed by finding
    a translation and scaling that aligns two tensors `A, B`
    defined as follows:
        ```
        T_A R_i + s_A T_i   = T_i'        ;  for all i    // · R_i^T
        s_A T_i R_i^T + T_A = T_i' R_i^T  ;  for all i
            ^^^^^^^^^         ^^^^^^^^^^
                A_i                B_i

        A_i := T_i R_i^T
        A = [A_1 A_2 ... A_N]
        B_i := T_i' R_i^T
        B = [B_1 B_2 ... B_N]
        ```
    The scale s_A can be retrieved by matching the correlations of
    the points sets A and B:
        ```
        s_A = (A-mean(A))*(B-mean(B)).sum() / ((A-mean(A))**2).sum()
        ```
    The translation `T_A` is then defined as:
        ```
        T_A = mean(B) - mean(A) * s_A
        ```
    """
    A = torch.bmm(cameras_src.R, cameras_src.T[:, :, None])[:, :, 0]
    B = torch.bmm(cameras_src.R, cameras_tgt.T[:, :, None])[:, :, 0]
    Amu = A.mean(0, keepdim=True)
    Bmu = B.mean(0, keepdim=True)
    if estimate_scale and A.shape[0] > 1:
        # get the scaling component by matching covariances
        # of centered A and centered B
        Ac = A - Amu
        Bc = B - Bmu
        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
        align_t_s = (Ac * Bc).mean() / (Ac**2).mean().clamp(eps)
    else:
        # set the scale to identity
        align_t_s = 1.0
    # get the translation as the difference between the means of A and B
    align_t_T = Bmu - align_t_s * Amu

    return align_t_R, align_t_T, align_t_s
//  [python]
def _align_camera_extrinsics(
    cameras_src: "CamerasBase",
    cameras_tgt: "CamerasBase",
    estimate_scale: bool = True,
    eps: float = 1e-9,
):  # pragma: no cover
    """
    Get the global rotation R_A with svd of cov(RR^T):
        
        R_A R_i = R_i' for all i
        R_A [R_1 R_2 ... R_N] = [R_1' R_2' ... R_N']
        U, _, V = svd([R_1 R_2 ... R_N]^T [R_1' R_2' ... R_N'])
        R_A = (U V^T)^T
        
    """
    RRcov = torch.bmm(cameras_src.R, cameras_tgt.R.transpose(2, 1)).mean(0)
    U, S, V = torch.linalg.svd(RRcov)
    align_t_R = V @ U.t()

    """
    The translation + scale `T_A` and `s_A` is computed by finding
    a translation and scaling that aligns two tensors `A, B`
    defined as follows:
        
        T_A R_i + s_A T_i   = T_i'        ;  for all i    // · R_i^T
        s_A T_i R_i^T + T_A = T_i' R_i^T  ;  for all i
            ^^^^^^^^^         ^^^^^^^^^^
                A_i                B_i

        A_i := T_i R_i^T
        A = [A_1 A_2 ... A_N]
        B_i := T_i' R_i^T
        B = [B_1 B_2 ... B_N]
        
    The scale s_A can be retrieved by matching the correlations of
    the points sets A and B:
        
        s_A = (A-mean(A))*(B-mean(B)).sum() / ((A-mean(A))**2).sum()
        
    The translation `T_A` is then defined as:
        
        T_A = mean(B) - mean(A) * s_A
        
    """
    A = torch.bmm(cameras_src.R, cameras_src.T[:, :, None])[:, :, 0]
    B = torch.bmm(cameras_src.R, cameras_tgt.T[:, :, None])[:, :, 0]
    Amu = A.mean(0, keepdim=True)
    Bmu = B.mean(0, keepdim=True)
    if estimate_scale and A.shape[0] > 1:
        # get the scaling component by matching covariances
        # of centered A and centered B
        Ac = A - Amu
        Bc = B - Bmu
        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
        align_t_s = (Ac * Bc).mean() / (Ac**2).mean().clamp(eps)
    else:
        # set the scale to identity
        align_t_s = 1.0
    # get the translation as the difference between the means of A and B
    align_t_T = Bmu - align_t_s * Amu

    return align_t_R, align_t_T, align_t_s
//  [python]
def zca(content, style):
    '''
    Matches the mean and covariance of 'content' to those of 'style'
    content -- N x D pytorch tensor of content feature vectors
    style   -- N x D pytorch tensor of style feature vectors
    '''
    mu_c = content.mean(0, keepdim=True)
    mu_s = style.mean(0, keepdim=True)

    content = content - mu_c
    style = style - mu_s

    cov_c = torch.matmul(content.transpose(1,0), content) / float(content.size(0))
    cov_s = torch.matmul(style.transpose(1,0), style) / float(style.size(0))

    u_c, sig_c, _ = torch.svd(cov_c + torch.eye(cov_c.size(0)).cuda()*1e-4)
    u_s, sig_s, _ = torch.svd(cov_s + torch.eye(cov_s.size(0)).cuda()*1e-4)

    sig_c = sig_c.unsqueeze(1)
    sig_s = sig_s.unsqueeze(1)


    u_c_i = u_c.transpose(1,0)
    u_s_i = u_s.transpose(1,0)

    scl_c = torch.sqrt(torch.clamp(sig_c, 1e-8, 1e8))
    scl_s = torch.sqrt(torch.clamp(sig_s, 1e-8, 1e8))


    whiten_c = whiten(content, u_c_i, u_c, scl_c)
    color_c = colorize(whiten_c, u_s_i, u_s, scl_s) + mu_s

    return color_c, cov_s