jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def mx_svd(mx, topk):
    """
    canonicalize the mx to target topk dim.
    Take SVD, take eigenvectors corresponding to smallest topk non zero evals, 
    Assuming input mx has 0 as eval, i.e. Laplacian mx.
    """
    U, D, V = torch.svd(mx) 
    #topk x m
    evecs = U.t()[-topk-1:-1]
    '''
    col1 = evecs[:, 0]
    idx = torch.argsort(col1)
    evecs = evecs[idx]
    '''
    row1 = evecs[-1, :]
    idx = torch.argsort(row1)
    evecs = evecs[:, idx]
    return evecs
//  [python]
def mx_svd(mx, topk):
    """
    canonicalize the mx to target topk dim.
    Take SVD, take eigenvectors corresponding to smallest topk non zero evals, 
    Assuming input mx has 0 as eval, i.e. Laplacian mx.
    """
    U, S, V = torch.linalg.svd(mx) 
    #topk x m
    evecs = U.t()[-topk-1:-1]
    '''
    col1 = evecs[:, 0]
    idx = torch.argsort(col1)
    evecs = evecs[idx]
    '''
    row1 = evecs[-1, :]
    idx = torch.argsort(row1)
    evecs = evecs[:, idx]
    return evecs
//  [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 symmetric_orthogonalization(x):
  """Maps 9D input vectors onto SO(3) via symmetric orthogonalization.

  x: should have size [batch_size, 9]

  Output has size [batch_size, 3, 3], where each inner 3x3 matrix is in SO(3).
  """
  m = x.view(-1, 3, 3)
  d = m.device
  u, s, v = torch.svd(m.cpu())
  u, v = u.to(d), v.to(d)
  vt = torch.transpose(v, 1, 2)
  det = torch.det(torch.bmm(u, vt))
  det = det.view(-1, 1, 1)
  vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
  r = torch.bmm(u, vt)
  return r
//  [python]
def symmetric_orthogonalization(x):
  """Maps 9D input vectors onto SO(3) via symmetric orthogonalization.

  x: should have size [batch_size, 9]

  Output has size [batch_size, 3, 3], where each inner 3x3 matrix is in SO(3).
  """
  m = x.view(-1, 3, 3)
  d = m.device
  u, s, v = torch.linalg.svd(m.cpu())
  u, v = u.to(d), v.to(d)
  vt = torch.transpose(v, 1, 2)
  det = torch.det(torch.bmm(u, vt))
  det = det.view(-1, 1, 1)
  vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
  r = torch.bmm(u, vt)
  return r
//  [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.

    help: https://gist.github.com/mkocabas/54ea2ff3b03260e3fedf8ad22536f427

    '''
    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, V = torch.svd(K)

    # 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)))      # position
    R = torch.linalg.inv(R)                 # rotation

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

    return R
//  [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.

    help: https://gist.github.com/mkocabas/54ea2ff3b03260e3fedf8ad22536f427

    '''
    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, V = torch.linalg.svd(K)

    # 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)))      # position
    R = torch.linalg.inv(R)                 # rotation

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

    return R
//  [python]
def estimate_transfo(xyz, xyz_target):
    """
    estimate the rotation and translation using Kabsch algorithm
    Parameters:
    xyz :
    xyz_target:
    """
    assert xyz.shape == xyz.shape
    xyz_c = xyz - xyz.mean(0)
    xyz_target_c = xyz_target - xyz_target.mean(0)
    Q = xyz_c.T.mm(xyz_target_c) / len(xyz)
    U, S, V = torch.svd(Q)
    d = torch.det(V.mm(U.T))
    diag = torch.diag(torch.tensor([1, 1, d], device=xyz.device))
    R = V.mm(diag).mm(U.T)
    t = xyz_target.mean(0) - R @ xyz.mean(0)
    T = torch.eye(4, device=xyz.device)
    T[:3, :3] = R
    T[:3, 3] = t
    return T
//  [python]
def estimate_transfo(xyz, xyz_target):
    """
    estimate the rotation and translation using Kabsch algorithm
    Parameters:
    xyz :
    xyz_target:
    """
    assert xyz.shape == xyz.shape
    xyz_c = xyz - xyz.mean(0)
    xyz_target_c = xyz_target - xyz_target.mean(0)
    Q = xyz_c.T.mm(xyz_target_c) / len(xyz)
    U, S, V = torch.linalg.svd(Q)
    d = torch.det(V.mm(U.T))
    diag = torch.diag(torch.tensor([1, 1, d], device=xyz.device))
    R = V.mm(diag).mm(U.T)
    t = xyz_target.mean(0) - R @ xyz.mean(0)
    T = torch.eye(4, device=xyz.device)
    T[:3, :3] = R
    T[:3, 3] = t
    return T
//  [python]
def subspace_initialize(self, data, hidden_dim, stacking_order=5):
        data_number, data_dim = data.shape[1], data.shape[0]
        hankel_matrix = hankel = scipy.linalg.hankel([x for x in range(stacking_order)],
                                                     [x for x in range(stacking_order - 1, data_dim, 1)])
        hankel_list = list(chain(*hankel_matrix.T.tolist()))
        expand_data = data[:, hankel_list]
        svd_data = expand_data.reshape((int(data_dim * stacking_order), -1))
        [u, s, v] = torch.svd(input=svd_data)
        s = torch.diag(s)
        big_c = u[:data_dim, :hidden_dim]
        big_z = s[:hidden_dim, :hidden_dim] @ v[:, :hidden_dim].transpose(1, 0)
        big_a = big_z[:, 1:] @ torch.linalg.pinv(big_z[:, 0:-1])
        return big_a, big_c, big_z
//  [python]
def subspace_initialize(self, data, hidden_dim, stacking_order=5):
        data_number, data_dim = data.shape[1], data.shape[0]
        hankel_matrix = hankel = scipy.linalg.hankel([x for x in range(stacking_order)],
                                                     [x for x in range(stacking_order - 1, data_dim, 1)])
        hankel_list = list(chain(*hankel_matrix.T.tolist()))
        expand_data = data[:, hankel_list]
        svd_data = expand_data.reshape((int(data_dim * stacking_order), -1))
        [u, s, v] = torch.linalg.svd(svd_data)
        s = torch.diag(s)
        big_c = u[:data_dim, :hidden_dim]
        big_z = s[:hidden_dim, :hidden_dim] @ v[:, :hidden_dim].transpose(1, 0)
        big_a = big_z[:, 1:] @ torch.linalg.pinv(big_z[:, 0:-1])
        return big_a, big_c, big_z
//  [python]
def compute_natural_gradient(self,g):
		
		uu,ss,vv = tr.svd(self.K.double())
		ss_inv,mask = self.pseudo_inverse(ss)
		ss_inv = tr.sqrt(ss_inv)
		vv = tr.einsum('i,ji->ij',ss_inv,vv)
		self.T = tr.einsum('ij,jk->ik', vv.float(), self.T)
		cond_g, G,D = self.make_system(g,mask)

		try:
			U  = tr.cholesky(G)
			cond_g = tr.cholesky_solve(cond_g.unsqueeze(-1), U).squeeze(-1)
		except:
			try:
				cond_g = tr.solve(cond_g.unsqueeze(-1), G)[0].squeeze(-1)
			except:
				pinv = tr.pinverse(G)
				cond_g = tr.einsum('mk,k',pinv,cond_g)
		cond_g = tr.einsum('md,m->d',self.T,cond_g)
		#cond_g = (g + self.beta * cond_g) / self.eps # grad(L) - M (sort of) # (1. - self.beta)
		cond_g = (g - self.beta * cond_g) / self.eps # original
		cond_g = D*cond_g
		return cond_g
//  [python]
def compute_natural_gradient(self,g):
		
		U, S, V = tr.linalg.svd(self.K.double())
		ss_inv,mask = self.pseudo_inverse(S)
		ss_inv = tr.sqrt(ss_inv)
		V = tr.einsum('i,ji->ij',ss_inv,V)
		self.T = tr.einsum('ij,jk->ik', V.float(), self.T)
		cond_g, G,D = self.make_system(g,mask)

		try:
			U  = tr.cholesky(G)
			cond_g = tr.cholesky_solve(cond_g.unsqueeze(-1), U).squeeze(-1)
		except:
			try:
				cond_g = tr.solve(cond_g.unsqueeze(-1), G)[0].squeeze(-1)
			except:
				pinv = tr.pinverse(G)
				cond_g = tr.einsum('mk,k',pinv,cond_g)
		cond_g = tr.einsum('md,m->d',self.T,cond_g)
		#cond_g = (g + self.beta * cond_g) / self.eps # grad(L) - M (sort of) # (1. - self.beta)
		cond_g = (g - self.beta * cond_g) / self.eps # original
		cond_g = D*cond_g
		return cond_g
//  [python]
def show_sv_for_increasing_aperture(conceptor, factor, title):
    """
    Show singular values for increasing aperture
    :param conceptors:
    :param factor:
    :param title:
    :return:
    """
    # Fig
    fig = plt.figure()
    ax = fig.gca()
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 1.5)
    ax.grid(True)

    # For each aperture multiplication
    for i in range(5):
        # Compute SVD
        _, S, _ = torch.svd(conceptor.get_C())

        # Plot
        ax.plot(S.numpy(), '--')

        # Multiply all conceptor's aperture by 10
        conceptor.multiply_aperture(factor)
    # end for

    # Show
    ax.set_xlabel(u"Singular values")
    ax.set_title(title)
    plt.show()
    plt.close()
//  [python]
def show_sv_for_increasing_aperture(conceptor, factor, title):
    """
    Show singular values for increasing aperture
    :param conceptors:
    :param factor:
    :param title:
    :return:
    """
    # Fig
    fig = plt.figure()
    ax = fig.gca()
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 1.5)
    ax.grid(True)

    # For each aperture multiplication
    for i in range(5):
        # Compute SVD
        _, S, _ = torch.linalg.svd(conceptor.get_C())

        # Plot
        ax.plot(S.numpy(), '--')

        # Multiply all conceptor's aperture by 10
        conceptor.multiply_aperture(factor)
    # end for

    # Show
    ax.set_xlabel(u"Singular values")
    ax.set_title(title)
    plt.show()
    plt.close()
//  [python]
def svd(self):
        """ 
        Efficient algorithm for finding Singular Value Decomposition, a tuple (U, S, Vh) for matrix M st S is a vector and U, Vh are orthogonal matrices, and U @ S.diag() @ Vh == M
        """
        Ua, Sa, Vha = torch.svd(self.A)
        Ub, Sb, Vhb = torch.svd(self.B)
        middle = Sa[..., :, None] * transpose(Vha) @ Ub * Sb[..., None, :]
        Um, Sm, Vhm = torch.svd(middle)
        U = Ua @ Um
        Vh = Vhb @ Vhm
        S = Sm
        return U, S, Vh
//  [python]
def svd(self):
        """ 
        Efficient algorithm for finding Singular Value Decomposition, a tuple (U, S, Vh) for matrix M st S is a vector and U, Vh are orthogonal matrices, and U @ S.diag() @ Vh == M
        """
        Ua, Sa, Vha = torch.linalg.svd(self.A)
        Ub, Sb, Vhb = torch.linalg.svd(self.B)
        middle = Sa[..., :, None] * transpose(Vha) @ Ub * Sb[..., None, :]
        Um, Sm, Vhm = torch.linalg.svd(middle)
        U = Ua @ Um
        Vh = Vhb @ Vhm
        S = Sm
        return U, S, Vh
//  [python]
def plot_singular_values(stats, title, xmin, xmax, ymin, ymax, log=False):
    """
    Plot singular values
    :param stats:
    :param title:
    :param timestep:
    :param start:
    :return:
    """
    # Compute R (correlation matrix)
    R = stats.t().mm(stats) / stats.shape[0]

    # Compute singular values
    U, S, V = torch.svd(R)
    singular_values = S

    # Compute singular values
    if log:
        singular_values = np.log10(singular_values)
    # end if

    # Fig
    fig = plt.figure()
    ax = fig.gca()
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.grid(True)

    # For each plot
    ax.plot(singular_values.numpy(), '--o')

    ax.set_xlabel("Timesteps")
    ax.set_title(title)
    plt.show()
    plt.close()

    return singular_values, U
//  [python]
def plot_singular_values(stats, title, xmin, xmax, ymin, ymax, log=False):
    """
    Plot singular values
    :param stats:
    :param title:
    :param timestep:
    :param start:
    :return:
    """
    # Compute R (correlation matrix)
    R = stats.t().mm(stats) / stats.shape[0]

    # Compute singular values
    U, S, V = torch.linalg.svd(R)
    singular_values = S

    # Compute singular values
    if log:
        singular_values = np.log10(singular_values)
    # end if

    # Fig
    fig = plt.figure()
    ax = fig.gca()
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.grid(True)

    # For each plot
    ax.plot(singular_values.numpy(), '--o')

    ax.set_xlabel("Timesteps")
    ax.set_title(title)
    plt.show()
    plt.close()

    return singular_values, U
//  [python]
def nuclear_pooling(tensor, kernel_size=None):
    shape = tensor.size()
    vectors = torch.zeros(1, shape[1], 1, 1).cuda()
    for i in range(shape[1]):
        u, s, v = torch.svd(tensor[0, i, :, :] + EPSILON)
        s_sum = torch.sum(s)
        vectors[0, i, 0, 0] = s_sum
    return vectors
//  [python]
def nuclear_pooling(tensor, kernel_size=None):
    shape = tensor.size()
    vectors = torch.zeros(1, shape[1], 1, 1).cuda()
    for i in range(shape[1]):
        u, s, v = torch.linalg.svd(tensor[0, i, :, :] + EPSILON)
        s_sum = torch.sum(s)
        vectors[0, i, 0, 0] = s_sum
    return vectors
//  [python]
def rot9d_to_rotmat(x):
    """
    Converts 9D rotation output to valid 3x3 rotation amtrix.
    Based on Levinson et al., An Analysis of SVD for Deep Rotation Estimation.

    Input:
        (B, 9)
    Output:
        (B, 9)
    """
    B = x.size()[0]
    x = x.reshape((B, 3, 3))
    u, s, v = torch.svd(x)

    v_T = v.transpose(-2, -1)
    s_p = torch.eye(3).to(x).reshape((1, 3, 3)).expand_as(x).clone()
    s_p[:, 2, 2] = torch.det(torch.matmul(u, v_T))
    x_out = torch.matmul(torch.matmul(u, s_p), v_T)

    return x_out.reshape((B, 9))
//  [python]
def rot9d_to_rotmat(x):
    """
    Converts 9D rotation output to valid 3x3 rotation amtrix.
    Based on Levinson et al., An Analysis of SVD for Deep Rotation Estimation.

    Input:
        (B, 9)
    Output:
        (B, 9)
    """
    B = x.size()[0]
    x = x.reshape((B, 3, 3))
    u, s, v = torch.linalg.svd(x)

    v_T = v.transpose(-2, -1)
    s_p = torch.eye(3).to(x).reshape((1, 3, 3)).expand_as(x).clone()
    s_p[:, 2, 2] = torch.det(torch.matmul(u, v_T))
    x_out = torch.matmul(torch.matmul(u, s_p), v_T)

    return x_out.reshape((B, 9))
//  [python]
def solve_overdetermined_linear_system(A, b, method='pinv'):
    # import pdb; pdb.set_trace()
    # A: [B, H, W, C]
    # b: [B, H, W]
    # method: 'pinv' or 'svd'
    assert method in ['pinv', 'svd']
    B = A.size(0)
    A = A.view(B, -1)
    b = b.view(B, -1)
    if method == 'pinv':
        return torch.pinverse(A) @ b
    else:
        u, s, v = torch.svd(A, some=False)
        s_new = torch.zeros(A.shape).to(device)
        for i in range(len(s)):
            s_new[i, i] = 1/s[i]
        s_inv = s_new.t()

        return v.t() @ s_inv @ u.t() @ b
//  [python]
def solve_overdetermined_linear_system(A, b, method='pinv'):
    assert method in ['pinv', 'svd']
    B = A.size(0)
    A = A.view(B, -1)
    b = b.view(B, -1)
    if method == 'pinv':
        return torch.pinverse(A) @ b
    else:
        u, s, v = torch.linalg.svd(A)
        s_new = torch.zeros(A.shape).to(device)
        for i in range(len(s)):
            s_new[i, i] = 1/s[i]
        s_inv = s_new.t()

        return v.t() @ s_inv @ u.t() @ b
//  [python]
def map_3d_to_2d(J3d, gt2d, gt3d):
    batch_size = J3d.shape[0]
    
    # Part 1: Align J3d with gt3d
    
    ### i. centering: must
    G3d = gt3d.clone()
    cent_J = torch.mean(J3d, dim=1, keepdim=True)
    J3d -= cent_J
    cent_G = torch.mean(G3d, dim=1, keepdim=True)
    G3d -= cent_G
    
    ### ii. scaling: not necessary here.
    
    ### iii. rotation
    M = torch.bmm(G3d.transpose(1,2), J3d) # [N, 3, 3]
    if platform == 'linux':
        U, D, V = batch_svd(M)
        R = torch.bmm(V, U.transpose(1,2))
    else:
        R = [None] * batch_size
        for i in range(batch_size):
            U, D, V = torch.svd(M[i])
            R[i] = torch.mm(V, U.transpose(0,1)) # transpose
        R = torch.stack(R, dim=0)
        
    reg3d = torch.bmm(J3d, R)
    ### eval stage I: rel error < 5%
    '''
    print(reg3d - G3d)
    test_case = range(10)
    for i in test_case:
        #np.savetxt('_test_cache/2d_gt_{}.xyz'.format(i), gt2d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_in_{}.xyz'.format(i), J3d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_reg_{}.xyz'.format(i), reg3d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_gt_{}.xyz'.format(i), G3d[i].detach().cpu().numpy(), delimiter=' ')
    '''
    
    # Part II: Align reg3d with gt2d
    G2d = gt2d.clone()
    reg2d = reg3d[:, :, :2]
    ### i. translation
    cent2d = torch.mean(G2d, dim=1, keepdim=True)
    G2d -= cent2d
    cent3d = torch.cat((cent2d,
        torch.zeros((batch_size, 1, 1), dtype=cent2d.dtype, device=cent2d.device)
    ), dim=2)
    
    ### ii. scaling
    G2d = G2d.view(batch_size, -1)
    r2d = reg2d.reshape(batch_size, -1)
    s = torch.sum(r2d * G2d, dim=1) / torch.sum(r2d * r2d, dim=1)
    s = s.unsqueeze(dim=1).unsqueeze(dim=2)
    
    ### eval: max 5 pixels error...
    '''
    reg2d = reg2d * s + cent2d
    print(reg2d - gt2d)
    # eval: make sure joint locations matching
    # 2D not very accurate, consider 3D?
    # gt3d and gt2d spatially aligned, roughly s * gt3d[:,:,:2] + t = gt2d
    
    gt2d = torch.cat((gt2d,
        torch.zeros((batch_size, gt2d.shape[1], 1), dtype=gt2d.dtype, device=gt2d.device)
    ), dim=2)
    reg2d = torch.cat((reg2d,
        torch.zeros((batch_size, reg2d.shape[1], 1), dtype=reg2d.dtype, device=reg2d.device)
    ), dim=2)
    
    test_case = range(10)
    for i in test_case:
        np.savetxt('_test_cache/2d_gt_{}.xyz'.format(i), gt2d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/2d_reg_{}.xyz'.format(i), reg2d[i].detach().cpu().numpy(), delimiter=' ')
    '''
    
    # Wrap-up the result into a function, keep z component as a depth inidcator
    def transform(x):
        return torch.bmm(x - cent_J, R) * s + cent3d
        
    return transform
//  [python]
def map_3d_to_2d(J3d, gt2d, gt3d):
    batch_size = J3d.shape[0]
    
    # Part 1: Align J3d with gt3d
    
    ### i. centering: must
    G3d = gt3d.clone()
    cent_J = torch.mean(J3d, dim=1, keepdim=True)
    J3d -= cent_J
    cent_G = torch.mean(G3d, dim=1, keepdim=True)
    G3d -= cent_G
    
    ### ii. scaling: not necessary here.
    
    ### iii. rotation
    M = torch.bmm(G3d.transpose(1,2), J3d) # [N, 3, 3]
    if platform == 'linux':
        U, S, V = torch.linalg.svd(M)
        R = torch.bmm(V, U.transpose(1,2))
    else:
        R = [None] * batch_size
        for i in range(batch_size):
            U, S, V = torch.linalg.svd(M[i])
            R[i] = torch.mm(V, U.transpose(0,1)) # transpose
        R = torch.stack(R, dim=0)
        
    reg3d = torch.bmm(J3d, R)
    ### eval stage I: rel error < 5%
    '''
    print(reg3d - G3d)
    test_case = range(10)
    for i in test_case:
        #np.savetxt('_test_cache/2d_gt_{}.xyz'.format(i), gt2d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_in_{}.xyz'.format(i), J3d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_reg_{}.xyz'.format(i), reg3d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/3d_gt_{}.xyz'.format(i), G3d[i].detach().cpu().numpy(), delimiter=' ')
    '''
    
    # Part II: Align reg3d with gt2d
    G2d = gt2d.clone()
    reg2d = reg3d[:, :, :2]
    ### i. translation
    cent2d = torch.mean(G2d, dim=1, keepdim=True)
    G2d -= cent2d
    cent3d = torch.cat((cent2d,
        torch.zeros((batch_size, 1, 1), dtype=cent2d.dtype, device=cent2d.device)
    ), dim=2)
    
    ### ii. scaling
    G2d = G2d.view(batch_size, -1)
    r2d = reg2d.reshape(batch_size, -1)
    s = torch.sum(r2d * G2d, dim=1) / torch.sum(r2d * r2d, dim=1)
    s = s.unsqueeze(dim=1).unsqueeze(dim=2)
    
    ### eval: max 5 pixels error...
    '''
    reg2d = reg2d * s + cent2d
    print(reg2d - gt2d)
    # eval: make sure joint locations matching
    # 2D not very accurate, consider 3D?
    # gt3d and gt2d spatially aligned, roughly s * gt3d[:,:,:2] + t = gt2d
    
    gt2d = torch.cat((gt2d,
        torch.zeros((batch_size, gt2d.shape[1], 1), dtype=gt2d.dtype, device=gt2d.device)
    ), dim=2)
    reg2d = torch.cat((reg2d,
        torch.zeros((batch_size, reg2d.shape[1], 1), dtype=reg2d.dtype, device=reg2d.device)
    ), dim=2)
    
    test_case = range(10)
    for i in test_case:
        np.savetxt('_test_cache/2d_gt_{}.xyz'.format(i), gt2d[i].detach().cpu().numpy(), delimiter=' ')
        np.savetxt('_test_cache/2d_reg_{}.xyz'.format(i), reg2d[i].detach().cpu().numpy(), delimiter=' ')
    '''
    
    # Wrap-up the result into a function, keep z component as a depth inidcator
    def transform(x):
        return torch.bmm(x - cent_J, R) * s + cent3d
        
    return transform
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, jsv_diff, condition_number = [], [], []
    if rand_data:
        picked_batch = np.random.permutation(len(data_loader))[:n_loop]
    else:
        picked_batch = list(range(n_loop))
    for i, (images, target) in enumerate(data_loader):
        if i in picked_batch:
            images, target = images.cuda(), target.cuda()
            batch_size = images.size(0)
            images.requires_grad = True # for Jacobian computation
            output = model(images)
            jacobian = compute_jacobian(images, output) # shape [batch_size, num_classes, num_channels, input_width, input_height]
            jacobian = jacobian.view(batch_size, num_classes, -1) # shape [batch_size, num_classes, num_channels*input_width*input_height]
            u, s, v = torch.svd(jacobian) # u: [batch_size, num_channels*input_width*input_height, num_classes], s: [batch_size, num_classes], v: [batch_size, num_channels*input_width*input_height, num_classes]
            s = s.data.cpu().numpy()
            jsv.append(s)
            jsv_diff.append((s - 1) ** 2)
            condition_number.append(s.max(axis=1) / s.min(axis=1))
            print_func('[%3d/%3d] calculating Jacobian...' % (i, len(data_loader)))
    jsv = np.concatenate(jsv)
    condition_number = np.concatenate(condition_number)
    return jsv, jsv_diff, condition_number
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, jsv_diff, condition_number = [], [], []
    if rand_data:
        picked_batch = np.random.permutation(len(data_loader))[:n_loop]
    else:
        picked_batch = list(range(n_loop))
    for i, (images, target) in enumerate(data_loader):
        if i in picked_batch:
            images, target = images.cuda(), target.cuda()
            batch_size = images.size(0)
            images.requires_grad = True # for Jacobian computation
            output = model(images)
            jacobian = compute_jacobian(images, output) # shape [batch_size, num_classes, num_channels, input_width, input_height]
            jacobian = jacobian.view(batch_size, num_classes, -1) # shape [batch_size, num_classes, num_channels*input_width*input_height]
            u, s, v = torch.linalg.svd(jacobian) # u: [batch_size, num_channels*input_width*input_height, num_classes], s: [batch_size, num_classes], v: [batch_size, num_channels*input_width*input_height, num_classes]
            s = s.data.cpu().numpy()
            jsv.append(s)
            jsv_diff.append((s - 1) ** 2)
            condition_number.append(s.max(axis=1) / s.min(axis=1))
            print_func('[%3d/%3d] calculating Jacobian...' % (i, len(data_loader)))
    jsv = np.concatenate(jsv)
    condition_number = np.concatenate(condition_number)
    return jsv, jsv_diff, condition_number
//  [python]
def estimate_R(self, pcl, return_quaternion=True):
        old_edges = torch.gather(input=self.init_pcl[:, None].repeat(1,self.K,1), dim=0, index=self.nn_idxs[..., None].repeat(1,1,3)) - self.init_pcl[:, None]  # N, K, 3
        edges = torch.gather(input=pcl[:, None].repeat(1,self.K,1), dim=0, index=self.nn_idxs[..., None].repeat(1,1,3)) - pcl[:, None]  # N, K, 3
        D = torch.diag_embed(self.weight, dim1=1, dim2=2)  # N, K, K
        S = torch.bmm(old_edges.permute(0, 2, 1), torch.bmm(D, edges))  # N, 3, 3
        unchanged = torch.unique(torch.where((edges == old_edges).all(dim=1))[0])
        S[unchanged] = 0
        U, _, W = torch.svd(S)
        R = torch.bmm(W, U.permute(0, 2, 1))
        if return_quaternion:
            q = matrix_to_quaternion(R)
            return q
        else:
            return R
//  [python]
def estimate_R(self, pcl, return_quaternion=True):
        old_edges = torch.gather(input=self.init_pcl[:, None].repeat(1,self.K,1), dim=0, index=self.nn_idxs[..., None].repeat(1,1,3)) - self.init_pcl[:, None]  # N, K, 3
        edges = torch.gather(input=pcl[:, None].repeat(1,self.K,1), dim=0, index=self.nn_idxs[..., None].repeat(1,1,3)) - pcl[:, None]  # N, K, 3
        D = torch.diag_embed(self.weight, dim1=1, dim2=2)  # N, K, K
        S = torch.bmm(old_edges.permute(0, 2, 1), torch.bmm(D, edges))  # N, 3, 3
        unchanged = torch.unique(torch.where((edges == old_edges).all(dim=1))[0])
        S[unchanged] = 0
        U, S, W = torch.linalg.svd(S)
        R = torch.bmm(W, U.permute(0, 2, 1))
        if return_quaternion:
            q = matrix_to_quaternion(R)
            return q
        else:
            return R
//  [python]
def horn87_v1(src, tgt, weight=None): # does not substract center, compare to horn87_np_v2
    '''
    # src: [(k), 3, n]
    # tgt: [(k), 3, n]
    # weight: [(k), n]
    # return: 
    #   (R, t) ([(k),3,3], [(k),3,1])
    '''
    if len(src.shape) == 2 and len(tgt.shape) == 2:
        src, tgt = src.unsqueeze(0), tgt.unsqueeze(0)
    assert(src.shape[2] == tgt.shape[2])
    nPts = src.shape[2]
    k = src.shape[0]
    has_weight=False
    if weight is None:
        weight = torch.ones(k,1,nPts).cuda().float()
    else:
        has_weight=True
        weight = weight.view(k,1,nPts)
    weight = weight / weight.sum(2,keepdim=True)
    src_ = src
    tgt_ = tgt
    if has_weight:
        for i in range(k):
            tgt_[i] *= weight[i]

    H = torch.bmm(src_, tgt_.transpose(2,1))
    R_ret = []
    for i in range(k):
        try:
            u, s, v = torch.svd(H[i,:,:].cpu())
            R = torch.matmul(v, u.t())
            det = torch.det(R)
            if det < 0:
                R = torch.matmul(v, torch.matmul(torch.diagflat(torch.FloatTensor([1,1,-1])),u.t()))
            R_ret.append(R.view(-1,3,3))

        except:
            print('rigid transform failed to converge')
            print('H:{}'.format(torch_op.npy(H)))

            R_ret.append(Variable(torch.eye(3).view(1,3,3), requires_grad=True))
    
    R_ret = torch.cat(R_ret).cuda()

    return R_ret
//  [python]
def horn87_v1(src, tgt, weight=None): # does not substract center, compare to horn87_np_v2
    '''
    # src: [(k), 3, n]
    # tgt: [(k), 3, n]
    # weight: [(k), n]
    # return: 
    #   (R, t) ([(k),3,3], [(k),3,1])
    '''
    if len(src.shape) == 2 and len(tgt.shape) == 2:
        src, tgt = src.unsqueeze(0), tgt.unsqueeze(0)
    assert(src.shape[2] == tgt.shape[2])
    nPts = src.shape[2]
    k = src.shape[0]
    has_weight=False
    if weight is None:
        weight = torch.ones(k,1,nPts).cuda().float()
    else:
        has_weight=True
        weight = weight.view(k,1,nPts)
    weight = weight / weight.sum(2,keepdim=True)
    src_ = src
    tgt_ = tgt
    if has_weight:
        for i in range(k):
            tgt_[i] *= weight[i]

    H = torch.bmm(src_, tgt_.transpose(2,1))
    R_ret = []
    for i in range(k):
        try:
            u, s, v = torch.linalg.svd(H[i,:,:].cpu())
            R = torch.matmul(v, u.t())
            det = torch.det(R)
            if det < 0:
                R = torch.matmul(v, torch.matmul(torch.diagflat(torch.FloatTensor([1,1,-1])),u.t()))
            R_ret.append(R.view(-1,3,3))

        except:
            print('rigid transform failed to converge')
            print('H:{}'.format(torch_op.npy(H)))

            R_ret.append(Variable(torch.eye(3).view(1,3,3), requires_grad=True))
    
    R_ret = torch.cat(R_ret).cuda()

    return R_ret
//  [python]
def wct(cont_feat, styl_feat):
    cFSize = cont_feat.size()
    c_mean = torch.mean(cont_feat, 1)
    c_mean = c_mean.unsqueeze(1).expand_as(cont_feat)
    cont_feat = cont_feat - c_mean
    iden = torch.eye(cFSize[0])
    iden = iden
    contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1] - 1)
    c_u, c_e, c_v = torch.svd(contentConv, some=False)
    k_c = cFSize[0]
    for i in range(cFSize[0] - 1, -1, -1):
        if c_e[i] >= 1e-05:
            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] - 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] >= 1e-05:
            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)
    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(cont_feat, styl_feat):
    cFSize = cont_feat.size()
    c_mean = torch.mean(cont_feat, 1)
    c_mean = c_mean.unsqueeze(1).expand_as(cont_feat)
    cont_feat = cont_feat - c_mean
    iden = torch.eye(cFSize[0])
    iden = iden
    contentConv = torch.mm(cont_feat, cont_feat.t()).div(cFSize[1] - 1)
    c_u, c_s, c_v = torch.linalg.svd(contentConv)
    k_c = cFSize[0]
    for i in range(cFSize[0] - 1, -1, -1):
        if c_s[i] >= 1e-05:
            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] - 1)
    s_u, s_s, s_v = torch.linalg.svd(styleConv)
    k_s = sFSize[0]
    for i in range(sFSize[0] - 1, -1, -1):
        if s_s[i] >= 1e-05:
            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_s[0:k_s].pow(0.5)
    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 kabsch_torch(X, Y, rot_mat=False):
    """ Kabsch alignment of X into Y.
        Assumes X,Y are both (D, N) - usually (3, N)
        * rot_mat: bool. whether to return rotation matrix.
    """
    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t())
    # Optimal rotation matrix via SVD - warning! W must be transposed
    if int(torch.__version__.split(".")[1]) < 8:
        V, S, W = torch.svd(C.detach())
        W = W.t()
    else:
        V, S, W = torch.linalg.svd(C.detach())
    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    if rot_mat:
        return X_, Y_, U
    return X_, Y_
//  [python]
def kabsch_torch(X, Y, rot_mat=False):
    """ Kabsch alignment of X into Y.
        Assumes X,Y are both (D, N) - usually (3, N)
        * rot_mat: bool. whether to return rotation matrix.
    """
    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t())
    # Optimal rotation matrix via SVD - warning! W must be transposed
    V, S, W = torch.linalg.svd(C.detach())
    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(W, V.t())
    # calculate rotations
    X_ = torch.matmul(U, X_).t()
    # return centered and aligned
    if rot_mat:
        return X_.t(), Y_, U
    return X_.t(), Y_
//  [python]
def guarded_matrix_solve_ls(A, b, W, condition_number_cap=1e5, sqrt_eps=1e-10, ls_l2_regularizer=1e-8):
    # Solve weighted least square ||\sqrt(W)(Ax-b)||^2
    # A - BxNxD
    # b - BxNx1
    # W - BxN
    batch_size, _, dim = A.size()
    sqrt_W = torch.sqrt(torch.clamp(W, min=sqrt_eps, max=None)).unsqueeze(2) # BxN
    A = A * sqrt_W # BxNxD
    b = b * sqrt_W # BxNx1
    # Compute singular value, trivializing the problem when condition number is too large
    AtA = torch.bmm(A.transpose(1,2), A)
    _, s, _ = torch.svd(AtA, compute_uv=False) # s will be BxD
    s = s.detach()
    mask = s[:,0] / s[:,-1] < condition_number_cap # B
    #import pdb; pdb.set_trace()
    #A = A * mask.float().view(batch_size, 1, 1)
    #x = torch.linalg.lstsq(A, b).solution
    AtA = AtA * mask.float().view(batch_size, 1, 1) + ls_l2_regularizer * torch.eye(dim).unsqueeze(0).to(A.device) # zero out badly conditioned data
    Atb = torch.bmm(A.transpose(1, 2) * mask.float().view(batch_size, 1, 1), b)
    x, _ = torch.solve(Atb, AtA)
    x = x.squeeze(2)
    return x
//  [python]
def guarded_matrix_solve_ls(A, b, W, condition_number_cap=1e5, sqrt_eps=1e-10, ls_l2_regularizer=1e-8):
    # Solve weighted least square ||\sqrt(W)(Ax-b)||^2
    # A - BxNxD
    # b - BxNx1
    # W - BxN
    batch_size, _, dim = A.size()
    sqrt_W = torch.sqrt(torch.clamp(W, min=sqrt_eps, max=None)).unsqueeze(2) # BxN
    A = A * sqrt_W # BxNxD
    b = b * sqrt_W # BxNx1
    # Compute singular value, trivializing the problem when condition number is too large
    AtA = torch.bmm(A.transpose(1,2), A)
    _, s, _ = torch.linalg.svd(AtA, compute_uv=False) # s will be BxD
    s = s.detach()
    mask = s[:,0] / s[:,-1] < condition_number_cap # B
    #import pdb; pdb.set_trace()
    #A = A * mask.float().view(batch_size, 1, 1)
    #x = torch.linalg.lstsq(A, b).solution
    AtA = AtA * mask.float().view(batch_size, 1, 1) + ls_l2_regularizer * torch.eye(dim).unsqueeze(0).to(A.device) # zero out badly conditioned data
    Atb = torch.bmm(A.transpose(1, 2) * mask.float().view(batch_size, 1, 1), b)
    x, _ = torch.solve(Atb, AtA)
    x = x.squeeze(2)
    return x
//  [python]
def decompose_model_paras(para_list1, para_list2, gradient_list, alpha):
    
    for i in range(len(para_list1)):
        
        para1 = para_list1[i]
        
        para2 = para_list2[i]
        
        grad = gradient_list[i]
        
        
        if len(para1.shape) <= 1:
            continue
        
        u1,s1,v1 = torch.svd(para1)
        
        u2,s2,v2 = torch.svd(para2)
        
        u0,s0,v0 = torch.svd(grad)
        
#         v = torch.mm(torch.inverse(torch.mm(torch.t(u1), u1) + torch.mm(torch.t(u2), u2)), torch.mm(torch.t(u1), para1) + torch.mm(torch.t(u2), para2))
        
        
        diff0 = torch.norm(u2*s2 - (u1*s1 - alpha*u0*s0))
        
        diff1 = torch.norm(s2*v2 - (s1*v1 - alpha*s0*v0))
        
        
        print(diff0, diff1)
//  [python]
def decompose_model_paras(para_list1, para_list2, gradient_list, alpha):
    
    for i in range(len(para_list1)):
        
        para1 = para_list1[i]
        
        para2 = para_list2[i]
        
        grad = gradient_list[i]
        
        
        if len(para1.shape) <= 1:
            continue
        
        u1,s1,v1 = torch.linalg.svd(para1)
        
        u2,s2,v2 = torch.linalg.svd(para2)
        
        u0,s0,v0 = torch.linalg.svd(grad)
        
#         v = torch.mm(torch.inverse(torch.mm(torch.t(u1), u1) + torch.mm(torch.t(u2), u2)), torch.mm(torch.t(u1), para1) + torch.mm(torch.t(u2), para2))
        
        
        diff0 = torch.norm(u2*s2 - (u1*s1 - alpha*u0*s0))
        
        diff1 = torch.norm(s2*v2 - (s1*v1 - alpha*s0*v0))
        
        
        print(diff0, diff1)
//  [python]
def normalize(cls, Rots):
        U, _, V = torch.svd(Rots)
        S = cls.Id.clone().repeat(Rots.shape[0], 1, 1)
        S[:, 2, 2] = torch.det(U) * torch.det(V)
        return U.bmm(S).bmm(V.transpose(1, 2))
//  [python]
def normalize(cls, Rots):
        U, S, V = torch.linalg.svd(Rots)
        S = cls.Id.clone().repeat(Rots.shape[0], 1, 1)
        S[:, 2, 2] = torch.det(U) * torch.det(V)
        return U.bmm(S).bmm(V.transpose(1, 2))
//  [python]
def _compute_sv_for_weight(weight, k_split, scheme):
        """Compute SVD for one layer."""
        # fold into matrix operator
        weight = scheme.fold(weight.detach())

        # get k_split as int instead of tensor
        k_split = k_split.item()

        # compute rank
        rank_k = min(weight.shape[0], weight.shape[1] // k_split)
        # pre-allocate singular values ...
        singular_values = torch.zeros(k_split, rank_k, device=weight.device)
        # compute singular values for each part of the decomposed weight
        for idx_k, w_per_g_k in enumerate(torch.chunk(weight, k_split, dim=1)):
            singular_values[idx_k] = torch.svd(w_per_g_k, compute_uv=False)[1]

        # note has shape [k_split x rank_per_k]
        return singular_values
//  [python]
def _compute_sv_for_weight(weight, k_split, scheme):
        """Compute SVD for one layer."""
        # fold into matrix operator
        weight = scheme.fold(weight.detach())

        # get k_split as int instead of tensor
        k_split = k_split.item()

        # compute rank
        rank_k = min(weight.shape[0], weight.shape[1] // k_split)
        # pre-allocate singular values ...
        singular_values = torch.zeros(k_split, rank_k, device=weight.device)
        # compute singular values for each part of the decomposed weight
        for idx_k, w_per_g_k in enumerate(torch.chunk(weight, k_split, dim=1)):
            singular_values[idx_k] = torch.linalg.svdvals(w_per_g_k)

        # note has shape [k_split x rank_per_k]
        return singular_values
//  [python]
def finish_pca(self):
        """Compute PCA from the data matrix."""
        # check that all data is valid or sub-sample
        idx_valid = self._data.sum(dim=0) != 0.0
        data = self._data[:, idx_valid]

        # center data and store mean
        data_mean = data.mean(dim=1, keepdim=True)
        data = data - data_mean
        self.data_mean = data_mean[:, 0]

        # covariance matrix
        cov_mat = data @ data.T

        # solve eigen-decomposition for "principle components"
        self.principle_components, s_values = torch.svd(cov_mat)[:2]

        # compute and store sensitivity
        # note that we have convention that "bigger is better"
        self.sensitivity = s_values / s_values.cumsum(0)
//  [python]
def finish_pca(self):
        """Compute PCA from the data matrix."""
        # check that all data is valid or sub-sample
        idx_valid = self._data.sum(dim=0) != 0.0
        data = self._data[:, idx_valid]

        # center data and store mean
        data_mean = data.mean(dim=1, keepdim=True)
        data = data - data_mean
        self.data_mean = data_mean[:, 0]

        # covariance matrix
        cov_mat = data @ data.T

        # solve eigen-decomposition for "principle components"
        self.principle_components, s_values = torch.linalg.svd(cov_mat)[:2]

        # compute and store sensitivity
        # note that we have convention that "bigger is better"
        self.sensitivity = s_values / s_values.cumsum(0)
//  [python]
def compute_similarity_transform(S1: torch.Tensor, S2: torch.Tensor) -> torch.Tensor:
    """
    Computes a similarity transform (sR, t) in a batched way that takes
    a set of 3D points S1 (B, N, 3) closest to a set of 3D points S2 (B, N, 3),
    where R is a 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    Args:
        S1 (torch.Tensor): First set of points of shape (B, N, 3).
        S2 (torch.Tensor): Second set of points of shape (B, N, 3).
    Returns:
        (torch.Tensor): The first set of points after applying the similarity transformation.
    """

    batch_size = S1.shape[0]
    S1 = S1.permute(0, 2, 1)
    S2 = S2.permute(0, 2, 1)
    # 1. Remove mean.
    mu1 = S1.mean(dim=2, keepdim=True)
    mu2 = S2.mean(dim=2, keepdim=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

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

    # 3. The outer product of X1 and X2.
    K = torch.matmul(X1, 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, V = torch.svd(K)
    Vh = V.permute(0, 2, 1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=U.device).unsqueeze(0).repeat(batch_size, 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.linalg.det(torch.matmul(U, Vh)))

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

    # 5. Recover scale.
    trace = torch.matmul(R, K).diagonal(offset=0, dim1=-1, dim2=-2).sum(dim=-1)
    scale = (trace / var1).unsqueeze(dim=-1).unsqueeze(dim=-1)

    # 6. Recover translation.
    t = mu2 - scale*torch.matmul(R, mu1)

    # 7. Error:
    S1_hat = scale*torch.matmul(R, S1) + t

    return S1_hat.permute(0, 2, 1)
//  [python]
def compute_similarity_transform(S1: torch.Tensor, S2: torch.Tensor) -> torch.Tensor:
    """
    Computes a similarity transform (sR, t) in a batched way that takes
    a set of 3D points S1 (B, N, 3) closest to a set of 3D points S2 (B, N, 3),
    where R is a 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    Args:
        S1 (torch.Tensor): First set of points of shape (B, N, 3).
        S2 (torch.Tensor): Second set of points of shape (B, N, 3).
    Returns:
        (torch.Tensor): The first set of points after applying the similarity transformation.
    """

    batch_size = S1.shape[0]
    S1 = S1.permute(0, 2, 1)
    S2 = S2.permute(0, 2, 1)
    # 1. Remove mean.
    mu1 = S1.mean(dim=2, keepdim=True)
    mu2 = S2.mean(dim=2, keepdim=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

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

    # 3. The outer product of X1 and X2.
    K = torch.matmul(X1, 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, V = torch.linalg.svd(K)
    Vh = V.permute(0, 2, 1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=U.device).unsqueeze(0).repeat(batch_size, 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.linalg.det(torch.matmul(U, Vh)))

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

    # 5. Recover scale.
    trace = torch.matmul(R, K).diagonal(offset=0, dim1=-1, dim2=-2).sum(dim=-1)
    scale = (trace / var1).unsqueeze(dim=-1).unsqueeze(dim=-1)

    # 6. Recover translation.
    t = mu2 - scale*torch.matmul(R, mu1)

    # 7. Error:
    S1_hat = scale*torch.matmul(R, S1) + t

    return S1_hat.permute(0, 2, 1)
//  [python]
def find_nobs(self, x_old, group_num):
        """
        :param x: group normalization [batch, feature dim]
        :return:
        """
        try:
            feature_dim, batch_size = x_old.size()
            assert feature_dim % group_num == 0
            x = x_old.view(group_num, feature_dim//group_num, batch_size)
            eps = 1e-4
            f_cov = (torch.bmm(x, x.transpose(1, 2)) / (batch_size - 1)).float()  # N * N
            eye = torch.eye(feature_dim//group_num).float().to(f_cov.device)
            x_stack = torch.FloatTensor().to(f_cov.device)
            for i in range(f_cov.size(0)):
                f_cov = torch.where(torch.isnan(f_cov), torch.zeros_like(f_cov), f_cov)
                U, S, V = torch.svd((1 - eps) * f_cov[i] + eps * eye)
                diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
                rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach()  # N * N
                x_transform = torch.mm(rotate_mtx, x[i])
                x_stack = torch.cat([x_stack, x_transform], dim=0)
            return x_stack
        except:
            return x_old
//  [python]
def find_nobs(self, x_old, group_num):
        """
        :param x: group normalization [batch, feature dim]
        :return:
        """
        try:
            feature_dim, batch_size = x_old.size()
            assert feature_dim % group_num == 0
            x = x_old.view(group_num, feature_dim//group_num, batch_size)
            eps = 1e-4
            f_cov = (torch.bmm(x, x.transpose(1, 2)) / (batch_size - 1)).float()  # N * N
            eye = torch.eye(feature_dim//group_num).float().to(f_cov.device)
            x_stack = torch.FloatTensor().to(f_cov.device)
            for i in range(f_cov.size(0)):
                f_cov = torch.where(torch.isnan(f_cov), torch.zeros_like(f_cov), f_cov)
                U, S, V = torch.linalg.svd((1 - eps) * f_cov[i] + eps * eye)
                diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
                rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach()  # N * N
                x_transform = torch.mm(rotate_mtx, x[i])
                x_stack = torch.cat([x_stack, x_transform], dim=0)
            return x_stack
        except:
            return x_old