jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _torch_svd_cast(input: Tensor) -> tuple[Tensor, Tensor, Tensor]:
	"""Helper function to make torch.svd work with other than fp32/64.

	The function torch.svd is only implemented for fp32/64 which makes
	impossible to be used by fp16 or others. What this function does, is cast
	input data type to fp32, apply torch.svd, and cast back to the input dtype.

	NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
	"""
	if not isinstance(input, Tensor):
		raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
	dtype = input.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32

	out1, out2, out3 = torch.svd(input.to(dtype))

	return out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype)
//  [python]
def _torch_svd_cast(input: Tensor) -> tuple[Tensor, Tensor, Tensor]:
	"""Helper function to make torch.svd work with other than fp32/64.

	The function torch.svd is only implemented for fp32/64 which makes
	impossible to be used by fp16 or others. What this function does, is cast
	input data type to fp32, apply torch.svd, and cast back to the input dtype.

	NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
	"""
	if not isinstance(input, Tensor):
		raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
	dtype = input.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32

	out1, out2, out3 = torch.linalg.svd(input.to(dtype))

	return out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype)
//  [python]
def RMSD(P, Q):
    '''Kabsch algorthm'''
    def rmsd(V, W):
        return torch.sqrt( torch.sum( (V-W)*(V-W) ) / len(V) )
    def centroid(X):
        return X.mean(axis=0)

    cP = centroid(P)
    cQ = centroid(Q)
    P = P - cP
    Q = Q - cQ

    # Computation of the covariance matrix
    C = torch.mm(P.T, Q)

    # Computate optimal rotation matrix using SVD
    V, S, W = torch.svd(C)

    # get sign( det(V)*det(W) ) to ensure right-handedness
    d = torch.ones([3,3],device=P.device)
    d[:,-1] = torch.sign(torch.det(V) * torch.det(W))

    # Rotation matrix U
    U = torch.mm(d*V, W.T)

    # Rotate P
    rP = torch.mm(P, U)

    # get RMS
    rms = rmsd(rP, Q)

    return rms
//  [python]
def RMSD(P, Q):
    '''Kabsch algorthm'''
    def rmsd(V, W):
        return torch.sqrt( torch.sum( (V-W)*(V-W) ) / len(V) )
    def centroid(X):
        return X.mean(axis=0)

    cP = centroid(P)
    cQ = centroid(Q)
    P = P - cP
    Q = Q - cQ

    # Computation of the covariance matrix
    C = torch.mm(P.T, Q)

    # Computate optimal rotation matrix using SVD
    U, S, V = torch.linalg.svd(C)

    # get sign( det(V)*det(W) ) to ensure right-handedness
    d = torch.ones([3,3],device=P.device)
    d[:,-1] = torch.sign(torch.det(U) * torch.det(V))

    # Rotation matrix U
    U = torch.mm(d*U, V.T)

    # Rotate P
    rP = torch.mm(P, U)

    # get RMS
    rms = rmsd(rP, Q)

    return rms
//  [python]
def wct(self, cont_feat, styl_feat):
    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] - 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] - 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)
    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(self, cont_feat, styl_feat):
    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] - 1) + iden
    # del iden
    c_u, c_s, c_v = torch.linalg.svd(contentConv)
    # 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] - 1)
    s_u, s_e, s_v = torch.linalg.svd(styleConv)

    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)
    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 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.
    '''
    if len(S1.shape) > 3:
        orig_shape = S1.shape
        S1 = S1.reshape(-1, *S1.shape[-2:])
        S2 = S2.reshape(-1, *S2.shape[-2:])
    else:
        orig_shape = None

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

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

    if orig_shape is not None:
         S1_hat = S1_hat.reshape(orig_shape)

    return S1_hat
//  [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.
    '''
    if len(S1.shape) > 3:
        orig_shape = S1.shape
        S1 = S1.reshape(-1, *S1.shape[-2:])
        S2 = S2.reshape(-1, *S2.shape[-2:])
    else:
        orig_shape = None

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

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

    if orig_shape is not None:
         S1_hat = S1_hat.reshape(orig_shape)

    return S1_hat
//  [python]
def compute_conceptor(X, aperture):
    """
    Compute conceptor
    :param X: Reservoir states (T x Nx)
    :param aperture: Aperture
    """
    x_length = X.size(0)
    x_dim = X.size(1)
    Rx = torch.mm(X.t(), X) / x_length
    Cx = torch.mm(torch.inverse(Rx + math.pow(aperture, -2) * torch.eye(x_dim)), Rx)
    URx, SRx, _ = torch.svd(Rx)
    UCx, SCx, _ = torch.svd(Cx)
    return Rx, Cx, URx, SRx, UCx, SCx
//  [python]
def compute_conceptor(X, aperture):
    """
    Compute conceptor
    :param X: Reservoir states (T x Nx)
    :param aperture: Aperture
    """
    x_length = X.size(0)
    x_dim = X.size(1)
    Rx = torch.mm(X.t(), X) / x_length
    Cx = torch.mm(torch.inverse(Rx + math.pow(aperture, -2) * torch.eye(x_dim)), Rx)
    URx, SRx, _ = torch.linalg.svd(Rx)
    UCx, SCx, _ = torch.linalg.svd(Cx)
    return Rx, Cx, URx, SRx, UCx, SCx
//  [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)
  u, s, v = torch.svd(m)
  vt = torch.transpose(v, 1, 2)
  det = torch.det(torch.matmul(u, vt))
  det = det.view(-1, 1, 1)
  vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
  r = torch.matmul(u, vt)
  return r


  '''Rot 6D representation'''
//  [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)
  u, s, v = torch.linalg.svd(m)
  vt = torch.transpose(v, 1, 2)
  det = torch.det(torch.matmul(u, vt))
  det = det.view(-1, 1, 1)
  vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
  r = torch.matmul(u, vt)
  return r
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        for i in range(input.shape[0]):
            U, S, V = torch.svd(input[i, :, :])
            Ss[i, :] = S
            Us[i, :, :] = U

        self.Us = Us
        self.Ss = Ss
        # self.save_for_backward(input)
        return Us, Ss
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        for i in range(input.shape[0]):
            U, S, V = torch.linalg.svd(input[i, :, :])
            Ss[i, :] = S
            Us[i, :, :] = U

        self.Us = Us
        self.Ss = Ss
        # self.save_for_backward(input)
        return Us, Ss
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        mSs = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.svd(input[i, :, :])

            Ss[i, :] = S
            Us[i, :, :] = U
            mSs[i, :, :] = torch.diag(torch.log(S))

        result = torch.matmul(Us, torch.matmul(mSs, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.save_for_backward(input)
        return result
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        mSs = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.linalg.svd(input[i, :, :])

            Ss[i, :] = S
            Us[i, :, :] = U
            mSs[i, :, :] = torch.diag(torch.log(S))

        result = torch.matmul(Us, torch.matmul(mSs, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.save_for_backward(input)
        return result
//  [python]
def svd_torch(x):
  x2 = torch.from_numpy(numpy.array(x))
  u, s, v = torch.svd(x2)
  u1 = u.numpy()
  s1 = s.numpy()
  v1 = v.numpy()
  return u1,s1,v1
//  [python]
def svd_torch(x):
  x2 = torch.from_numpy(numpy.array(x))
  u, s, v = torch.linalg.svd(x2)
  u1 = u.numpy()
  s1 = s.numpy()
  v1 = v.numpy()
  return u1,s1,v1
//  [python]
def procrustes_torch_parallel(p_gt, p_pred):
    # p_gt and p_pred need to be of shape (-1, 3, #joints)
    # care: run on cpu! way faster than on gpu

    mu_gt = p_gt.mean(dim=2)
    mu_pred = p_pred.mean(dim=2)

    X0 = p_gt - mu_gt[:, :, None]
    Y0 = p_pred - mu_pred[:, :, None]

    ssX = (X0**2.).sum(dim=(1, 2))
    ssY = (Y0**2.).sum(dim=(1, 2))

    # centred Frobenius norm
    normX = torch.sqrt(ssX) + 1e-6
    normY = torch.sqrt(ssY) + 1e-6

    # scale to equal (unit) norm
    X0 /= normX[:, None, None]
    Y0 /= normY[:, None, None]

    # optimum rotation matrix of Y
    A = torch.bmm(X0, Y0.transpose(1, 2))

    try:
        U, s, V = torch.svd(A, some=True, compute_uv=True)
    except:
        print("ERROR IN SVD, could not converge")
        print("SVD INPUT IS:")
        print(A)
        print(A.shape)
        print('A NaNs:', A.isnan().sum())
        print('GT NaNs', X0.isnan().sum())
        print('Pred NaNs', Y0.isnan().sum())
        print('normX', normX.min(), normX.max())
        print('normY', normY.min(), normY.max())
        return None

    T = torch.bmm(V, U.transpose(1, 2))

    # Make sure we have a rotation
    detT = torch.det(T)
    sign = torch.sign(detT)
    V[:, :, -1] *= sign[:, None]
    s[:, -1] *= sign
    T = torch.bmm(V, U.transpose(1, 2))

    traceTA = s.sum(dim=1)

    # optimum scaling of Y
    b = traceTA * normX / normY

    # standardised distance between X and b*Y*T + c
    d = 1 - traceTA**2

    # transformed coords
    scale = normX*traceTA
    Z = (scale[:, None, None] * torch.bmm(Y0.transpose(1, 2), T) + mu_gt[:, None, :]).transpose(1, 2)

    # transformation matrix
    c = mu_gt - b[:, None]*(torch.bmm(mu_pred[:, None, :], T)).squeeze()

    # transformation values
    tform = {'rotation': T, 'scale': b, 'translation': c}
    return d, Z, tform
//  [python]
def procrustes_torch_parallel(p_gt, p_pred):
    # p_gt and p_pred need to be of shape (-1, 3, #joints)
    # care: run on cpu! way faster than on gpu

    mu_gt = p_gt.mean(dim=2)
    mu_pred = p_pred.mean(dim=2)

    X0 = p_gt - mu_gt[:, :, None]
    Y0 = p_pred - mu_pred[:, :, None]

    ssX = (X0**2.).sum(dim=(1, 2))
    ssY = (Y0**2.).sum(dim=(1, 2))

    # centred Frobenius norm
    normX = torch.sqrt(ssX) + 1e-6
    normY = torch.sqrt(ssY) + 1e-6

    # scale to equal (unit) norm
    X0 /= normX[:, None, None]
    Y0 /= normY[:, None, None]

    # optimum rotation matrix of Y
    A = torch.bmm(X0, Y0.transpose(1, 2))

    try:
        U, s, V = torch.linalg.svd(A)
    except:
        print("ERROR IN SVD, could not converge")
        print("SVD INPUT IS:")
        print(A)
        print(A.shape)
        print('A NaNs:', A.isnan().sum())
        print('GT NaNs', X0.isnan().sum())
        print('Pred NaNs', Y0.isnan().sum())
        print('normX', normX.min(), normX.max())
        print('normY', normY.min(), normY.max())
        return None

    T = torch.bmm(V, U.transpose(1, 2))

    # Make sure we have a rotation
    detT = torch.det(T)
    sign = torch.sign(detT)
    V[:, :, -1] *= sign[:, None]
    s[:, -1] *= sign
    T = torch.bmm(V, U.transpose(1, 2))

    traceTA = s.sum(dim=1)

    # optimum scaling of Y
    b = traceTA * normX / normY

    # standardised distance between X and b*Y*T + c
    d = 1 - traceTA**2

    # transformed coords
    scale = normX*traceTA
    Z = (scale[:, None, None] * torch.bmm(Y0.transpose(1, 2), T) + mu_gt[:, None, :]).transpose(1, 2)

    # transformation matrix
    c = mu_gt - b[:, None]*(torch.bmm(mu_pred[:, None, :], T)).squeeze()

    # transformation values
    tform = {'rotation': T, 'scale': b, 'translation': c}
    return d, Z, tform
//  [python]
def solve_svd(A, Y):
    U, S, V = torch.svd(A)
    return V @ ((U.t() @ Y) / S)
//  [python]
def solve_svd(A, Y):
    U, S, V = torch.linalg.svd(A)
    return V @ ((U.t() @ Y) / S)
//  [python]
def rigid_tranform_from_points(data, choice, device):
    # data:[n,c]
    # choice:[min_sample]

    P_p = data[choice, 0:3]
    P = data[choice, 3:6]
    #print(P_p.shape, P.shape)
    P_p_mean = torch.mean(P_p, dim=0, keepdim=True) # 1,3
    P_mean = torch.mean(P, dim=0, keepdim=True) # 1,3
    Q_p = P_p - P_p_mean # n_matches,3
    Q = P - P_mean # n_matches,3

    A = torch.matmul(Q.transpose(-1,-2), Q_p)# 3,3
    #B = A.cpu()
    u, s, v = torch.svd(A) # 3,3
    reflect = torch.eye(3, device=device)
    R = torch.matmul(u, v.transpose(-1,-2)) # 3， 3
    reflect[2, 2] = torch.det(R)
    R = torch.matmul(u, reflect).matmul(v.transpose(-1,-2)) # 3， 3
    t = (P_mean.transpose(-1,-2) - torch.matmul(R, P_p_mean.transpose(-1,-2))).squeeze(-1) # 3, 3 * 3, 1 -> 3, 1 -> 3
    
    I = torch.eye(4, device=device).view(4, 4)
    I[:3, :3] = R
    I[:3, 3] = t
    h = I.view(-1)# 16
    if torch.isnan(h).sum():
        print(R, t)
        sys.exit(-1)
    return h
//  [python]
def rigid_tranform_from_points(data, choice, device):
    # data:[n,c]
    # choice:[min_sample]

    P_p = data[choice, 0:3]
    P = data[choice, 3:6]
    #print(P_p.shape, P.shape)
    P_p_mean = torch.mean(P_p, dim=0, keepdim=True) # 1,3
    P_mean = torch.mean(P, dim=0, keepdim=True) # 1,3
    Q_p = P_p - P_p_mean # n_matches,3
    Q = P - P_mean # n_matches,3

    A = torch.matmul(Q.transpose(-1,-2), Q_p)# 3,3
    #B = A.cpu()
    u, s, v = torch.linalg.svd(A) # 3,3
    reflect = torch.eye(3, device=device)
    R = torch.matmul(u, v.transpose(-1,-2)) # 3， 3
    reflect[2, 2] = torch.det(R)
    R = torch.matmul(u, reflect).matmul(v.transpose(-1,-2)) # 3， 3
    t = (P_mean.transpose(-1,-2) - torch.matmul(R, P_p_mean.transpose(-1,-2))).squeeze(-1) # 3, 3 * 3, 1 -> 3, 1 -> 3
    
    I = torch.eye(4, device=device).view(4, 4)
    I[:3, :3] = R
    I[:3, 3] = t
    h = I.view(-1)# 16
    if torch.isnan(h).sum():
        print(R, t)
        sys.exit(-1)
    return h
//  [python]
def haar_measure(out):
    """
    samples a k-dimensional subspace in R^d distributed according to the Haar measure on the Grassmannian Gr(k, d)
    out is assumed to be a tensor of dimension (bs, d, k)
    """

    bs, d, k = out.shape 
    z = torch.randn(bs, k, d) / sqrt(2)
    for i in range(bs):
        _, _, out[i, :, :] = torch.svd(z[i, :, :], some = True)
//  [python]
def haar_measure(out):
    """
    samples a k-dimensional subspace in R^d distributed according to the Haar measure on the Grassmannian Gr(k, d)
    out is assumed to be a tensor of dimension (bs, d, k)
    """

    bs, d, k = out.shape 
    z = torch.randn(bs, k, d) / sqrt(2)
    for i in range(bs):
        _, _, out[i, :, :] = torch.linalg.svd(z[i, :, :])
//  [python]
def svd(x):
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    batches = x.shape[:-2]
    if batches:
        n, m = x.shape[-2:]
        k = min(n, m)
        U, d, V = x.new(*batches, n, k), x.new(*batches, k), x.new(*batches, m, k)
        for idx in itertools.product(*map(range, batches)):
            U[idx], d[idx], V[idx] = torch.svd(x[idx])
        return U, d, V
    else:
        return torch.svd(x)
//  [python]
def svd(x):
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    batches = x.shape[:-2]
    if batches:
        n, m = x.shape[-2:]
        k = min(n, m)
        U, d, V = x.new(*batches, n, k), x.new(*batches, k), x.new(*batches, m, k)
        for idx in itertools.product(*map(range, batches)):
            U[idx], d[idx], V[idx] = torch.linalg.svd(x[idx])
        return U, d, V
    else:
        return torch.linalg.svd(x)
//  [python]
def _so_maxpool(self, x):
        while len(x.data.shape) < 4:
            x = torch.unsqueeze(x, 0)
        batchSize, tupleLen, nFeat, dimFeat = x.data.shape
        x = torch.reshape(x, (-1, dimFeat))
        x = torch.unsqueeze(x, -1)
        x = x.matmul(x.transpose(1, 2))

        x = torch.reshape(x, (batchSize, tupleLen, nFeat, dimFeat, dimFeat))
        x = torch.max(x, 2).values
        x = torch.reshape(x, (-1, dimFeat, dimFeat))
        if self.do_pe:
            x = x.double()
            # u_, s_, vh_ = torch.linalg.svd(x)
            # dist = torch.dist(x, u_ @ torch.diag_embed(s_) @ vh_)
            # dist_same = torch.allclose(x, u_ @ torch.diag_embed(s_) @ vh_)
            # s_alpha = torch.pow(s_, 0.5)
            # x = u_ @ torch.diag_embed(s_alpha) @ vh_

            # For pytorch versions < 1.9
            u_, s_, v_ = torch.svd(x)
            # dist = torch.dist(x, u_ @ torch.diag_embed(s_) @  v_.transpose(-2, -1))
            # dist_same = torch.allclose(x, u_ @ torch.diag_embed(s_) @  v_.transpose(-2, -1))
            s_alpha = torch.pow(s_, 0.5)
            x = u_ @ torch.diag_embed(s_alpha) @ v_.transpose(-2, -1)

        x = torch.reshape(x, (batchSize, tupleLen, dimFeat, dimFeat))
        return x
//  [python]
def _so_maxpool(self, x):
        while len(x.data.shape) < 4:
            x = torch.unsqueeze(x, 0)
        batchSize, tupleLen, nFeat, dimFeat = x.data.shape
        x = torch.reshape(x, (-1, dimFeat))
        x = torch.unsqueeze(x, -1)
        x = x.matmul(x.transpose(1, 2))

        x = torch.reshape(x, (batchSize, tupleLen, nFeat, dimFeat, dimFeat))
        x = torch.max(x, 2).values
        x = torch.reshape(x, (-1, dimFeat, dimFeat))
        if self.do_pe:
            x = x.double()
            u_, s_, vh_ = torch.linalg.svd(x)
            s_alpha = torch.pow(s_, 0.5)
            x = u_ @ torch.diag_embed(s_alpha) @ vh_

        x = torch.reshape(x, (batchSize, tupleLen, dimFeat, dimFeat))
        return x
//  [python]
def zca(x):
    sigma = torch.mm(x.t(), x) / x.shape[0]
    U, S, _ = torch.svd(sigma)
    pcs = torch.mm(torch.mm(U, torch.diag(1. / torch.sqrt(S + 1e-7))), U.t())

    # Apply ZCA whitening
    whitex = torch.mm(x, pcs)
    return whitex
//  [python]
def zca(x):
    sigma = torch.mm(x.t(), x) / x.shape[0]
    U, S, _ = torch.linalg.svd(sigma)
    pcs = torch.mm(torch.mm(U, torch.diag(1. / torch.sqrt(S + 1e-7))), U.t())

    # Apply ZCA whitening
    whitex = torch.mm(x, pcs)
    return whitex
//  [python]
def _orthogonalize_columns(matrix):
    matrix -= torch.mean(matrix, 1).unsqueeze(1)
    D, V = torch.svd(matrix)[1:]
    W = torch.matmul(V, torch.diag((1./(torch.sqrt(D) + 1e-10))))
    return torch.matmul(matrix, W)
//  [python]
def _orthogonalize_columns(matrix):
    matrix -= torch.mean(matrix, 1).unsqueeze(1)
    U, S, V = torch.linalg.svd(matrix)
    W = torch.matmul(V, torch.diag((1./(torch.sqrt(S) + 1e-10))))
    return torch.matmul(matrix, W)
//  [python]
def set_SVD(self,centers):
        U_c, S_c, V_c = torch.svd(centers, some=False, compute_uv=True)
        self.U = U_c
        self.s = S_c
        self.V = V_c
        #resclae s
        self.s = self.base_value*torch.abs(S_c/self.base_value)**self.sharp_factor
        self.s = torch.sign(S_c)*self.s
//  [python]
def set_SVD(self,centers):
        U_c, S_c, V_c = torch.linalg.svd(centers, full_matrices=False)
        self.U = U_c
        self.s = S_c
        self.V = V_c
        #resclae s
        self.s = self.base_value*torch.abs(S_c/self.base_value)**self.sharp_factor
        self.s = torch.sign(S_c)*self.s
//  [python]
print(torch.asin(r))

# ...and linear algebra operations like determinant and singular value decomposition
print('\nDeterminant of r:')
print(torch.det(r))
print('\nSingular value decomposition of r:')
print(torch.svd(r))
//  [python]
print(torch.linalg.svd(r))
//  [python]
def pca(X, k, center=True):
    """
    Principal Components Analysis impl. with SVD function

    :param X:
    :param k:
    :param center:
    :return:
    """

    n = X.size()[0]
    ones = torch.ones(n).view([n, 1])
    h = ((1 / n) * torch.mm(ones, ones.t())) if center else torch.zeros(n * n).view([n, n])
    H = torch.eye(n) - h
    X_center = torch.mm(H.double(), X.double())
    u, s, v = torch.svd(X_center)
    components = v[:k].t()
    explained_variance = torch.mul(s[:k], s[:k]) / (n - 1)
    return {'X': X, 'k': k, 'components': components,
            'explained_variance': explained_variance}
//  [python]
def pca(X, k, center=True):
    """
    Principal Components Analysis impl. with SVD function

    :param X:
    :param k:
    :param center:
    :return:
    """

    n = X.size()[0]
    ones = torch.ones(n).view([n, 1])
    h = ((1 / n) * torch.mm(ones, ones.t())) if center else torch.zeros(n * n).view([n, n])
    H = torch.eye(n) - h
    X_center = torch.mm(H.double(), X.double())
    u, s, v = torch.linalg.svd(X_center)
    components = v[:k].t()
    explained_variance = torch.mul(s[:k], s[:k]) / (n - 1)
    return {'X': X, 'k': k, 'components': components,
            'explained_variance': explained_variance}
//  [python]
def U_UBi_Shb(self, Vs, vs):

        # compute U and V
        V = torch.cat([torch.sqrt(vs[i]) * V for i, V in enumerate(Vs)], 1)
        U = V / torch.sqrt(vs[-1])
        eye = torch.eye(U.shape[1]).cuda()
        B = torch.mm(torch.transpose(U, 0, 1), U) + eye
        # cholB = torch.potrf(B, upper=False)
        # Bi = torch.potri(cholB, upper=False)
        Ub, Shb, Vb = torch.svd(B)
        # Bi = (Vb / Shb).mm(torch.transpose(Vb, 0, 1))
        Bi = torch.inverse(B)
        UBi = torch.mm(U, Bi)

        return U, UBi, Shb
//  [python]
def U_UBi_Shb(self, Vs, vs):

        # compute U and V
        V = torch.cat([torch.sqrt(vs[i]) * V for i, V in enumerate(Vs)], 1)
        U = V / torch.sqrt(vs[-1])
        eye = torch.eye(U.shape[1]).cuda()
        B = torch.mm(torch.transpose(U, 0, 1), U) + eye
        # cholB = torch.potrf(B, upper=False)
        # Bi = torch.potri(cholB, upper=False)
        Ub, Shb, Vb = torch.linalg.svd(B)
        # Bi = (Vb / Shb).mm(torch.transpose(Vb, 0, 1))
        Bi = torch.inverse(B)
        UBi = torch.mm(U, Bi)

        return U, UBi, Shb
//  [python]
def nll_ineff(self, X, Vs):
        vs = self.get_vs()
        V = torch.cat([torch.sqrt(vs[i]) * V for i, V in enumerate(Vs)], 1)
        K = V.mm(V.transpose(0, 1)) + vs[-1] * torch.eye(X.shape[0]).cuda()
        Uk, Shk, Vk = torch.svd(K)
        Ki = torch.inverse(K)
        # Ki = (Vk / Shk).mm(torch.transpose(Vk, 0, 1))
        Xb = Ki.mm(X)

        quad_term = torch.einsum("ij,ij->i", [X, Xb])[:, None]
        logdetK = X.shape[1] * torch.log(Shk).sum()
        nll = 0.5 * quad_term + 0.5 * logdetK / X.shape[0]

        return nll
//  [python]
def nll_ineff(self, X, Vs):
        vs = self.get_vs()
        V = torch.cat([torch.sqrt(vs[i]) * V for i, V in enumerate(Vs)], 1)
        K = V.mm(V.transpose(0, 1)) + vs[-1] * torch.eye(X.shape[0]).cuda()
        Uk, Shk, Vk = torch.linalg.svd(K)
        Ki = torch.inverse(K)
        # Ki = (Vk / Shk).mm(torch.transpose(Vk, 0, 1))
        Xb = Ki.mm(X)

        quad_term = torch.einsum("ij,ij->i", [X, Xb])[:, None]
        logdetK = X.shape[1] * torch.log(Shk).sum()
        nll = 0.5 * quad_term + 0.5 * logdetK / X.shape[0]

        return nll
//  [python]
def __init__(self, gan, samples_per_row=4, session=None):
        BaseSampler.__init__(self, gan, samples_per_row)
        self.latent1 = self.gan.latent.next()
        self.latent2 = self.gan.latent.next()
        self.velocity = 2/30.0
        direction = self.gan.latent.next()
        self.origin = direction
        self.pos = self.latent1
        self.hardtanh = nn.Hardtanh()
        g_params = self.gan.latent_parameters()
        if self.latent1.shape[1] // 2 == g_params[0].shape[1]:
            #recombine a split
            g_params = [torch.cat([p1, p2], 1) for p1, p2 in zip(g_params[:len(g_params)//2], g_params[len(g_params)//2:])]
            
        self.eigvec = torch.svd(torch.cat(g_params, 0)).V
        #self.eigvec = torch.svd(list(self.gan.g_parameters())[0]).V
        self.index = 0
        self.direction = self.eigvec[:, self.index].unsqueeze(0)
        self.direction = self.direction / torch.norm(self.direction)
        self.ones = torch.ones_like(self.direction, device="cuda:0")
        self.mask = torch.cat([torch.zeros([1, direction.shape[1]//2]), torch.ones([1, direction.shape[1]//2])], dim=1).cuda()
        self.mask = torch.ones_like(self.mask).cuda()
        self.steps = 30
//  [python]
def __init__(self, gan, samples_per_row=4, session=None):
        BaseSampler.__init__(self, gan, samples_per_row)
        self.latent1 = self.gan.latent.next()
        self.latent2 = self.gan.latent.next()
        self.velocity = 2/30.0
        direction = self.gan.latent.next()
        self.origin = direction
        self.pos = self.latent1
        self.hardtanh = nn.Hardtanh()
        g_params = self.gan.latent_parameters()
        if self.latent1.shape[1] // 2 == g_params[0].shape[1]:
            #recombine a split
            g_params = [torch.cat([p1, p2], 1) for p1, p2 in zip(g_params[:len(g_params)//2], g_params[len(g_params)//2:])]
            
        self.eigvec = torch.linalg.svd(torch.cat(g_params, 0)).V
        #self.eigvec = torch.linalg.svd(list(self.gan.g_parameters())[0]).V
        self.index = 0
        self.direction = self.eigvec[:, self.index].unsqueeze(0)
        self.direction = self.direction / torch.norm(self.direction)
        self.ones = torch.ones_like(self.direction, device="cuda:0")
        self.mask = torch.cat([torch.zeros([1, direction.shape[1]//2]), torch.ones([1, direction.shape[1]//2])], dim=1).cuda()
        self.mask = torch.ones_like(self.mask).cuda()
        self.steps = 30
//  [python]
def _matrix_power(matrix: torch.Tensor, power: float) -> torch.Tensor:
    # use CPU for svd for speed up
    device = matrix.device
    matrix = matrix.cpu()
    u, s, v = torch.svd(matrix)
    return (u @ s.pow_(power).diag() @ v.t()).to(device)
//  [python]
def _matrix_power(matrix: torch.Tensor, power: float) -> torch.Tensor:
    # use CPU for svd for speed up
    device = matrix.device
    matrix = matrix.cpu()
    u, s, v = torch.linalg.svd(matrix)
    return (u @ s.pow_(power).diag() @ v.t()).to(device)
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, 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)
            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, condition_number
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, 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)
            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, condition_number
//  [python]
def forward(self, input, some=True, compute_uv=True, out=None):
        """
        Params
        ------
        input: torch.Tensor
            Shape: (bs, n, m) or (n, m)
        some: bool
            If some is True (default), the method returns the reduced singular value decomposition
            i.e., if the last two dimensions of input are m and n, then the returned U and V matrices
            will contain only min(n, m)min(n,m) orthonormal columns.
        compute_uv: bool
            If compute_uv is False, the returned U and V matrices will be zero matrices of shape
            m \times m)(m×m) and (n \times n)(n×n) respectively. some will be ignored here.
        out:
        Returns
        -------
        U, S, V
        """
        return torch.svd(input, some=some, compute_uv=compute_uv, out=out)
//  [python]
def forward(self, input, some=True, compute_uv=True, out=None):
        """
        Params
        ------
        input: torch.Tensor
            Shape: (bs, n, m) or (n, m)
        some: bool
            If some is True (default), the method returns the reduced singular value decomposition
            i.e., if the last two dimensions of input are m and n, then the returned U and V matrices
            will contain only min(n, m)min(n,m) orthonormal columns.
        compute_uv: bool
            If compute_uv is False, the returned U and V matrices will be zero matrices of shape
            m \times m)(m×m) and (n \times n)(n×n) respectively. some will be ignored here.
        out:
        Returns
        -------
        U, S, V
        """
        return torch.linalg.svd(input, some=some, compute_uv=compute_uv, out=out)
//  [python]
def align_coords_Kabsch(p_cycle_coords, q_cycle_coords, p_mask, q_mask=None):
    """
    align p_cycle_coords with q_cycle_coords

    mask indicates which atoms to apply RMSD minimization over; these atoms are used to calculate the
    final rotation and translation matrices, which are applied to ALL atoms
    """
    if not q_mask:
        q_mask = p_mask

    q_cycle_coords_centered = q_cycle_coords[:, q_mask] - q_cycle_coords[:, q_mask].mean(dim=1, keepdim=True)
    p_cycle_coords_centered = p_cycle_coords[:, :, p_mask] - p_cycle_coords[:, :, p_mask].mean(dim=2, keepdim=True)

    H = torch.matmul(p_cycle_coords_centered.permute(0, 1, 3, 2), q_cycle_coords_centered.unsqueeze(0))
    u, s, v = torch.svd(H)
    d = torch.sign(torch.det(torch.matmul(v, u.permute(0, 1, 3, 2))))
    R_1 = torch.diag_embed(torch.ones([p_cycle_coords.size(0), q_cycle_coords.size(0), 3]))
    R_1[:, :, 2, 2] = d
    R = torch.matmul(v, torch.matmul(R_1, u.permute(0, 1, 3, 2)))
    b = q_cycle_coords[:, q_mask].mean(dim=1) - torch.matmul(R, p_cycle_coords[:, :, p_mask].mean(dim=2).unsqueeze(
        -1)).squeeze(-1)

    p_cycle_coords_aligned = torch.matmul(R, p_cycle_coords.permute(0, 1, 3, 2)).permute(0, 1, 3, 2) + b.unsqueeze(2)

    return p_cycle_coords_aligned
//  [python]
def align_coords_Kabsch(p_cycle_coords, q_cycle_coords, p_mask, q_mask=None):
    """
    align p_cycle_coords with q_cycle_coords

    mask indicates which atoms to apply RMSD minimization over; these atoms are used to calculate the
    final rotation and translation matrices, which are applied to ALL atoms
    """
    if not q_mask:
        q_mask = p_mask

    q_cycle_coords_centered = q_cycle_coords[:, q_mask] - q_cycle_coords[:, q_mask].mean(dim=1, keepdim=True)
    p_cycle_coords_centered = p_cycle_coords[:, :, p_mask] - p_cycle_coords[:, :, p_mask].mean(dim=2, keepdim=True)

    H = torch.matmul(p_cycle_coords_centered.permute(0, 1, 3, 2), q_cycle_coords_centered.unsqueeze(0))
    u, s, v = torch.linalg.svd(H)
    d = torch.sign(torch.det(torch.matmul(v, u.permute(0, 1, 3, 2))))
    R_1 = torch.diag_embed(torch.ones([p_cycle_coords.size(0), q_cycle_coords.size(0), 3]))
    R_1[:, :, 2, 2] = d
    R = torch.matmul(v, torch.matmul(R_1, u.permute(0, 1, 3, 2)))
    b = q_cycle_coords[:, q_mask].mean(dim=1) - torch.matmul(R, p_cycle_coords[:, :, p_mask].mean(dim=2).unsqueeze(
        -1)).squeeze(-1)

    p_cycle_coords_aligned = torch.matmul(R, p_cycle_coords.permute(0, 1, 3, 2)).permute(0, 1, 3, 2) + b.unsqueeze(2)

    return p_cycle_coords_aligned
//  [python]
def decorrelate_data(data):
    """flattened = tensor.flatten(2).squeeze(0) # this code aint shit, yo
    cov_matrix = torch.cov(flattened)
    sqrt_inv_cov_matrix = torch.linalg.inv(torch.sqrt(cov_matrix))
    decorrelated_tensor = torch.dot(flattened, sqrt_inv_cov_matrix.T)
    decorrelated_tensor = decorrelated_tensor.unflatten(2, tensor.shape[2:]).unsqueeze(0)"""

    # Reshape the 4D tensor to a 2D tensor for covariance calculation
    num_samples, num_channels, height, width = data.size()
    data_reshaped = data.view(num_samples, num_channels, -1)
    data_reshaped = data_reshaped - torch.mean(data_reshaped, dim=2, keepdim=True)

    # Compute covariance matrix
    cov_matrix = torch.matmul(data_reshaped, data_reshaped.transpose(1, 2)) / (height * width - 1)

    # Compute the inverse square root of the covariance matrix
    u, s, v = torch.svd(cov_matrix)
    sqrt_inv_cov_matrix = torch.matmul(u, torch.matmul(torch.diag_embed(1.0 / torch.sqrt(s)), v.transpose(1, 2)))

    # Reshape sqrt_inv_cov_matrix to match the dimensions of data_reshaped
    sqrt_inv_cov_matrix = sqrt_inv_cov_matrix.unsqueeze(0).expand(num_samples, -1, -1, -1)

    # Decorrelate the data
    decorrelated_data = torch.matmul(data_reshaped.transpose(1, 2), sqrt_inv_cov_matrix.transpose(2, 3))
    decorrelated_data = decorrelated_data.transpose(2, 3)
    
    # Reshape back to the original shape
    decorrelated_data = decorrelated_data.view(num_samples, num_channels, height, width)

    return decorrelated_data.to(data.device)
//  [python]
def decorrelate_data(data):
    """flattened = tensor.flatten(2).squeeze(0) # this code aint shit, yo
    cov_matrix = torch.cov(flattened)
    sqrt_inv_cov_matrix = torch.linalg.inv(torch.sqrt(cov_matrix))
    decorrelated_tensor = torch.dot(flattened, sqrt_inv_cov_matrix.T)
    decorrelated_tensor = decorrelated_tensor.unflatten(2, tensor.shape[2:]).unsqueeze(0)"""

    # Reshape the 4D tensor to a 2D tensor for covariance calculation
    num_samples, num_channels, height, width = data.size()
    data_reshaped = data.view(num_samples, num_channels, -1)
    data_reshaped = data_reshaped - torch.mean(data_reshaped, dim=2, keepdim=True)

    # Compute covariance matrix
    cov_matrix = torch.matmul(data_reshaped, data_reshaped.transpose(1, 2)) / (height * width - 1)

    # Compute the inverse square root of the covariance matrix
    u, s, v = torch.linalg.svd(cov_matrix)
    sqrt_inv_cov_matrix = torch.matmul(u, torch.matmul(torch.diag_embed(1.0 / torch.sqrt(s)), v.transpose(1, 2)))

    # Reshape sqrt_inv_cov_matrix to match the dimensions of data_reshaped
    sqrt_inv_cov_matrix = sqrt_inv_cov_matrix.unsqueeze(0).expand(num_samples, -1, -1, -1)

    # Decorrelate the data
    decorrelated_data = torch.matmul(data_reshaped.transpose(1, 2), sqrt_inv_cov_matrix.transpose(2, 3))
    decorrelated_data = decorrelated_data.transpose(2, 3)
    
    # Reshape back to the original shape
    decorrelated_data = decorrelated_data.view(num_samples, num_channels, height, width)

    return decorrelated_data.to(data.device)
//  [python]
def get_matrix_kernel(A, eps = 1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

    :param A: matrix
    :return: matrix where each row is a basis vector of the kernel of A
    '''
    _u, s, v = torch.svd(A)
    kernel = v.t()[s < eps]
    return kernel
//  [python]
def get_matrix_kernel(A, eps = 1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

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

    S1 = S1.to(torch.float32)
    S2 = S2.to(torch.float32)

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

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

    S1 = S1.to(torch.float32)
    S2 = S2.to(torch.float32)

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

    # 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
//  [python]
def match_colors_for_image_set(image_set, style_img):
    """
    image_set: [N, H, W, 3]
    style_img: [H, W, 3]
    """
    sh = image_set.shape
    image_set = image_set.view(-1, 3)
    style_img = style_img.view(-1, 3).to(image_set.device)

    mu_c = image_set.mean(0, keepdim=True)
    mu_s = style_img.mean(0, keepdim=True)

    cov_c = torch.matmul((image_set - mu_c).transpose(1, 0), image_set - mu_c) / float(image_set.size(0))
    cov_s = torch.matmul((style_img - mu_s).transpose(1, 0), style_img - mu_s) / float(style_img.size(0))

    u_c, sig_c, _ = torch.svd(cov_c)
    u_s, sig_s, _ = torch.svd(cov_s)

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

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

    tmp_mat = u_s @ scl_s @ u_s_i @ u_c @ scl_c @ u_c_i
    tmp_vec = mu_s.view(1, 3) - mu_c.view(1, 3) @ tmp_mat.T

    image_set = image_set @ tmp_mat.T + tmp_vec.view(1, 3)
    image_set = image_set.contiguous().clamp_(0.0, 1.0).view(sh)

    color_tf = torch.eye(4).float().to(tmp_mat.device)
    color_tf[:3, :3] = tmp_mat
    color_tf[:3, 3:4] = tmp_vec.T
    return image_set, color_tf
//  [python]
def match_colors_for_image_set(image_set, style_img):
    """
    image_set: [N, H, W, 3]
    style_img: [H, W, 3]
    """
    sh = image_set.shape
    image_set = image_set.view(-1, 3)
    style_img = style_img.view(-1, 3).to(image_set.device)

    mu_c = image_set.mean(0, keepdim=True)
    mu_s = style_img.mean(0, keepdim=True)

    cov_c = torch.matmul((image_set - mu_c).transpose(1, 0), image_set - mu_c) / float(image_set.size(0))
    cov_s = torch.matmul((style_img - mu_s).transpose(1, 0), style_img - mu_s) / float(style_img.size(0))

    u_c, sig_c, _ = torch.linalg.svd(cov_c)
    u_s, sig_s, _ = torch.linalg.svd(cov_s)

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

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

    tmp_mat = u_s @ scl_s @ u_s_i @ u_c @ scl_c @ u_c_i
    tmp_vec = mu_s.view(1, 3) - mu_c.view(1, 3) @ tmp_mat.T

    image_set = image_set @ tmp_mat.T + tmp_vec.view(1, 3)
    image_set = image_set.contiguous().clamp_(0.0, 1.0).view(sh)

    color_tf = torch.eye(4).float().to(tmp_mat.device)
    color_tf[:3, :3] = tmp_mat
    color_tf[:3, 3:4] = tmp_vec.T
    return image_set, color_tf
//  [python]
def project_matrix_to_unitary(W):
    if isinstance(W, np.ndarray):
        U, _, V = np.linalg.svd(W, full_matrices=True)
        U_refine = np.matmul(U, V)
    elif isinstance(W, torch.Tensor):
        U, _, V = torch.svd(W, some=False)
        U_refine = torch.matmul(U, V.transpose(-2, -1))
    else:
        raise NotImplementedError
    return U_refine