jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        rmsds = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            lig_coords_pred_mean = lig_coords_pred.mean(dim=0, keepdim=True)  # (1,3)
            lig_coords_mean = lig_coords.mean(dim=0, keepdim=True)  # (1,3)

            A = (lig_coords_pred - lig_coords_pred_mean).transpose(0, 1) @ (lig_coords - lig_coords_mean)

            U, S, Vt = torch.svd(A)

            corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=lig_coords_pred.device))
            rotation = (U @ corr_mat) @ Vt
            translation = lig_coords_pred_mean - torch.t(rotation @ lig_coords_mean.t())  # (1,3)

            lig_coords = (rotation @ lig_coords.t()).t() + translation
            rmsds.append(torch.sqrt(torch.mean(torch.sum(((lig_coords_pred - lig_coords) ** 2), dim=1))))
        return torch.tensor(rmsds).mean()
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        rmsds = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            lig_coords_pred_mean = lig_coords_pred.mean(dim=0, keepdim=True)  # (1,3)
            lig_coords_mean = lig_coords.mean(dim=0, keepdim=True)  # (1,3)

            A = (lig_coords_pred - lig_coords_pred_mean).transpose(0, 1) @ (lig_coords - lig_coords_mean)

            U, S, Vt = torch.linalg.svd(A)

            corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=lig_coords_pred.device))
            rotation = (U @ corr_mat) @ Vt
            translation = lig_coords_pred_mean - torch.t(rotation @ lig_coords_mean.t())  # (1,3)

            lig_coords = (rotation @ lig_coords.t()).t() + translation
            rmsds.append(torch.sqrt(torch.mean(torch.sum(((lig_coords_pred - lig_coords) ** 2), dim=1))))
        return torch.tensor(rmsds).mean()
//  [python]
def orthogonalize_left2right(self, nt, way, dc=-1,
                                 normalize=False):
        # dc=-1意味着不进行裁剪
        assert nt < len(self.tensors)-1
        s = self.tensors[nt].shape
        if 0 < dc < s[-1]:
            # In this case, truncation is required
            way = 'svd'
            if_trun = True
        else:
            if_trun = False

        tensor = self.tensors[nt].reshape(-1, s[-1]).to('cpu')
        if way.lower() == 'svd':
            u, s, v = tc.svd(tensor)
            lm = s.to(dtype=u.dtype)
            if if_trun:
                u = u[:, :dc].to(device=self.device)
                r = tc.diag(lm[:dc]).to(device=self.device).mm(
                    v[:dc, :].to(device=self.device))
            else:
                u = u.to(device=self.device)
                r = tc.diag(lm).to(device=self.device).mm(v.to(device=self.device))
        else:
            u, r = tc.linalg.qr(tensor)
            lm = None
            u, r = u.to(device=self.device), r.to(device=self.device)
        self.tensors[nt] = u.reshape(s[0], s[1], -1)
        if normalize:
            r /= tc.norm(r)
        self.tensors[nt+1] = tc.tensordot(
            r, self.tensors[nt+1], [[1], [0]])
        return lm
//  [python]
def orthogonalize_left2right(self, nt, way, dc=-1,
                                 normalize=False):
        # dc=-1意味着不进行裁剪
        assert nt < len(self.tensors)-1
        s = self.tensors[nt].shape
        if 0 < dc < s[-1]:
            # In this case, truncation is required
            way = 'svd'
            if_trun = True
        else:
            if_trun = False

        tensor = self.tensors[nt].reshape(-1, s[-1]).to('cpu')
        if way.lower() == 'svd':
            u, lm, v = tc.linalg.svd(tensor,
                                     full_matrices=False)
            lm = lm.to(dtype=u.dtype)
            if if_trun:
                u = u[:, :dc].to(device=self.device)
                r = tc.diag(lm[:dc]).to(device=self.device).mm(
                    v[:dc, :].to(device=self.device))
            else:
                u = u.to(device=self.device)
                r = tc.diag(lm).to(device=self.device).mm(v.to(device=self.device))
        else:
            u, r = tc.linalg.qr(tensor)
            lm = None
            u, r = u.to(device=self.device), r.to(device=self.device)
        self.tensors[nt] = u.reshape(s[0], s[1], -1)
        if normalize:
            r /= tc.norm(r)
        self.tensors[nt+1] = tc.tensordot(
            r, self.tensors[nt+1], [[1], [0]])
        return lm
//  [python]
def orthogonalize_right2left(self, nt, way, dc=-1, normalize=False):
        # dc=-1意味着不进行裁剪
        assert nt > 0
        s = self.tensors[nt].shape
        if 0 < dc < s[0]:
            # In this case, truncation is required
            way = 'svd'
            if_trun = True
        else:
            if_trun = False

        tensor = self.tensors[nt].reshape(s[0], -1).t().to('cpu')
        if way.lower() == 'svd':
            u, s, v = tc.svd(tensor)
            lm = s.to(dtype=u.dtype)
            if if_trun:
                u = u[:, :dc].to(device=self.device)
                r = tc.diag(lm[:dc]).to(device=self.device).mm(v[:dc, :].to(self.device))
            else:
                u = u.to(device=self.device)
                r = tc.diag(lm).to(device=self.device).mm(v.to(device=self.device))
        else:
            u, r = tc.qr(tensor)
            lm = None
            u, r = u.to(device=self.device), r.to(device=self.device)
        self.tensors[nt] = u.t().reshape(-1, s[1], s[2])
        if normalize:
            r /= tc.norm(r)
        self.tensors[nt - 1] = tc.tensordot(self.tensors[nt - 1], r, [[2], [1]])
        return lm
//  [python]
def orthogonalize_right2left(self, nt, way, dc=-1, normalize=False):
        # dc=-1意味着不进行裁剪
        assert nt > 0
        s = self.tensors[nt].shape
        if 0 < dc < s[0]:
            # In this case, truncation is required
            way = 'svd'
            if_trun = True
        else:
            if_trun = False

        tensor = self.tensors[nt].reshape(s[0], -1).t().to('cpu')
        if way.lower() == 'svd':
            u, lm, v = tc.linalg.svd(tensor, full_matrices=False)
            lm = lm.to(dtype=u.dtype)
            if if_trun:
                u = u[:, :dc].to(device=self.device)
                r = tc.diag(lm[:dc]).to(device=self.device).mm(v[:dc, :].to(self.device))
            else:
                u = u.to(device=self.device)
                r = tc.diag(lm).to(device=self.device).mm(v.to(device=self.device))
        else:
            u, r = tc.linalg.qr(tensor)
            lm = None
            u, r = u.to(device=self.device), r.to(device=self.device)
        self.tensors[nt] = u.t().reshape(-1, s[1], s[2])
        if normalize:
            r /= tc.norm(r)
        self.tensors[nt - 1] = tc.tensordot(self.tensors[nt - 1], r, [[2], [1]])
        return lm
//  [python]
def __init__(
        self, m, img_shape, fast=True, device="cpu", dtype=torch.float32, **kwargs
    ):
        super().__init__(**kwargs)
        self.name = f"spcamera_m{m}"
        self.img_shape = img_shape
        self.fast = fast
        self.device = device

        if self.fast:
            C, H, W = img_shape
            mi = min(int(np.sqrt(m)), H)
            mj = min(m - mi, W)

            revi = get_permutation_list(H)[:mi]
            revj = get_permutation_list(W)[:mj]

            assert H == 1 << int(np.log2(H)), "image height must be a power of 2"
            assert W == 1 << int(np.log2(W)), "image width must be a power of 2"

            mask = torch.zeros(img_shape).unsqueeze(0)
            for i in range(len(revi)):
                for j in range(len(revj)):
                    mask[0, :, revi[i], revj[j]] = 1

            mask = mask.to(device)

        else:
            n = int(np.prod(img_shape[1:]))
            A = torch.ones((m, n), device=device)
            A[torch.randn_like(A) > 0.5] = -1.0
            A /= np.sqrt(m)  # normalize
            u, s, vh = torch.svd(A)

            mask = s.to(device).unsqueeze(0).type(dtype)
            self.vh = vh.to(device).type(dtype)
            self.u = u.to(device).type(dtype)

            self.u = torch.nn.Parameter(self.u, requires_grad=False)
            self.vh = torch.nn.Parameter(self.vh, requires_grad=False)

        self.mask = torch.nn.Parameter(mask, requires_grad=False)
//  [python]
def __init__(
        self, m, img_shape, fast=True, device="cpu", dtype=torch.float32, **kwargs
    ):
        super().__init__(**kwargs)
        self.name = f"spcamera_m{m}"
        self.img_shape = img_shape
        self.fast = fast
        self.device = device

        if self.fast:
            C, H, W = img_shape
            mi = min(int(np.sqrt(m)), H)
            mj = min(m - mi, W)

            revi = get_permutation_list(H)[:mi]
            revj = get_permutation_list(W)[:mj]

            assert H == 1 << int(np.log2(H)), "image height must be a power of 2"
            assert W == 1 << int(np.log2(W)), "image width must be a power of 2"

            mask = torch.zeros(img_shape).unsqueeze(0)
            for i in range(len(revi)):
                for j in range(len(revj)):
                    mask[0, :, revi[i], revj[j]] = 1

            mask = mask.to(device)

        else:
            n = int(np.prod(img_shape[1:]))
            A = torch.ones((m, n), device=device)
            A[torch.randn_like(A) > 0.5] = -1.0
            A /= np.sqrt(m)  # normalize
            u, mask, vh = torch.linalg.svd(A, full_matrices=False)

            mask = mask.to(device).unsqueeze(0).type(dtype)
            self.vh = vh.to(device).type(dtype)
            self.u = u.to(device).type(dtype)

            self.u = torch.nn.Parameter(self.u, requires_grad=False)
            self.vh = torch.nn.Parameter(self.vh, requires_grad=False)

        self.mask = torch.nn.Parameter(mask, requires_grad=False)
//  [python]
def kabsch_rotation(P, Q):
    """
    Using the Kabsch algorithm with two sets of paired point P and Q, centered
    around the centroid. Each vector set is represented as an NxD
    matrix, where D is the the dimension of the space.
    The algorithm works in three steps:
    - a centroid translation of P and Q (assumed done before this function
      call)
    - the computation of a covariance matrix C
    - computation of the optimal rotation matrix U
    For more info see http://en.wikipedia.org/wiki/Kabsch_algorithm
    Parameters
    ----------
    P : array
        (N,D) matrix, where N is points and D is dimension.
    Q : array
        (N,D) matrix, where N is points and D is dimension.
    Returns
    -------
    U : matrix
        Rotation matrix (D,D)
    """

    # Computation of the covariance matrix
    C = P.transpose(-1, -2) @ Q

    # Computation of the optimal rotation matrix
    # This can be done using singular value decomposition (SVD)
    # Getting the sign of the det(V)*(W) to decide
    # whether we need to correct our rotation matrix to ensure a
    # right-handed coordinate system.
    # And finally calculating the optimal rotation matrix U
    # see http://en.wikipedia.org/wiki/Kabsch_algorithm
    U, S, V = torch.svd(C)
    d = (torch.det(U) * torch.det(V)) < 0.0

    if d:
        U[:, -1] = -U[:, -1]

    # Create Rotation matrix U
    U = U @ V
    return U
//  [python]
def kabsch_rotation(P, Q):
    """
    Using the Kabsch algorithm with two sets of paired point P and Q, centered
    around the centroid. Each vector set is represented as an NxD
    matrix, where D is the the dimension of the space.
    The algorithm works in three steps:
    - a centroid translation of P and Q (assumed done before this function
      call)
    - the computation of a covariance matrix C
    - computation of the optimal rotation matrix U
    For more info see http://en.wikipedia.org/wiki/Kabsch_algorithm
    Parameters
    ----------
    P : array
        (N,D) matrix, where N is points and D is dimension.
    Q : array
        (N,D) matrix, where N is points and D is dimension.
    Returns
    -------
    U : matrix
        Rotation matrix (D,D)
    """

    # Computation of the covariance matrix
    C = P.transpose(-1, -2) @ Q

    # Computation of the optimal rotation matrix
    # This can be done using singular value decomposition (SVD)
    # Getting the sign of the det(V)*(W) to decide
    # whether we need to correct our rotation matrix to ensure a
    # right-handed coordinate system.
    # And finally calculating the optimal rotation matrix U
    # see http://en.wikipedia.org/wiki/Kabsch_algorithm
    V, _, W = torch.linalg.svd(C)
    d = (torch.linalg.det(V) * torch.linalg.det(W)) < 0.0

    if d:
        V[:, -1] = -V[:, -1]

    # Create Rotation matrix U
    U = V @ W
    return U
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        #Logit of Block 4 feature
        feat1 = model.module.intermediate_forward(inputs,layer_index=4)
        B, C, H, W = feat1.size()
        feat1 = feat1.view(B, C, H * W)
        u, s, v = torch.svd(feat1)
        feat1 = feat1 - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat1 = feat1 - power_iteration(feat1, iter=20)
        feat1 = feat1.view(B,C,H,W)
        logits1 = model.module.head(model.module.before_head(feat1))

        # Logit of Block 4 feature
        feat2 = model.module.intermediate_forward(inputs, layer_index=3)
        B, C, H, W = feat2.size()
        feat2 = feat2.view(B, C, H * W)
        u, s, v = torch.svd(feat2)
        feat2 = feat2 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat2 = feat2 - power_iteration(feat2, iter=20)
        feat2 = feat2.view(B, C, H, W)
        feat2 = model.module.body.block4(feat2)
        logits2 = model.module.head(model.module.before_head(feat2))

        #Fusion at the logit space
        logits = (logits1+logits2) / 2
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        #Logit of Block 4 feature
        feat1 = model.module.intermediate_forward(inputs,layer_index=4)
        B, C, H, W = feat1.size()
        feat1 = feat1.view(B, C, H * W)
        u,s,v = torch.linalg.svd(feat1,full_matrices=False)
        feat1 = feat1 - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat1 = feat1 - power_iteration(feat1, iter=20)
        feat1 = feat1.view(B,C,H,W)
        logits1 = model.module.head(model.module.before_head(feat1))

        # Logit of Block 4 feature
        feat2 = model.module.intermediate_forward(inputs, layer_index=3)
        B, C, H, W = feat2.size()
        feat2 = feat2.view(B, C, H * W)
        u, s, v = torch.linalg.svd(feat2,full_matrices=False)
        feat2 = feat2 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat2 = feat2 - power_iteration(feat2, iter=20)
        feat2 = feat2.view(B, C, H, W)
        feat2 = model.module.body.block4(feat2)
        logits2 = model.module.head(model.module.before_head(feat2))

        #Fusion at the logit space
        logits = (logits1+logits2) / 2
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.module.body.block4[-1].conv3.weight.data
    B, C, H, W = weight.size()
    weight = weight.squeeze()
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.body.block4[-1].conv3.weight.data = weight

    weight = model.module.body.block4[-1].conv2.weight.data
    B, C, H, W = weight.size()
    weight = weight.view(B, C * H * W)
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.body.block4[-1].conv2.weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        feat = model.module.intermediate_forward(inputs,layer_index=4)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.svd(feat)
        feat = feat - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B,C,H,W)
        logits = model.module.head(model.module.before_head(feat))
        
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    
    weight = model.module.body.block4[-1].conv3.weight.data
    B, C, H, W = weight.size()
    weight = weight.squeeze()
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.body.block4[-1].conv3.weight.data = weight

    weight = model.module.body.block4[-1].conv2.weight.data
    B, C, H, W = weight.size()
    weight = weight.view(B, C * H * W)
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.body.block4[-1].conv2.weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        feat = model.module.intermediate_forward(inputs,layer_index=4)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u,s,v = torch.linalg.svd(feat,full_matrices=False)
        feat = feat - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        #if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        #feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B,C,H,W)
        logits = model.module.head(model.module.before_head(feat))
        
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def get_optimal_transform_v2(
    p: torch.Tensor,
    q: torch.Tensor,
    m: torch.Tensor,
    num_dim: int = 1,
    eps: float = 1e-6,
) -> torch.Tensor:
    """
    calculate u such that p @ u ~ q.
    p, q has shape  [*, *dim, 3]
    mask has shape  [*, *dim]
    ret has shape     [*, *dim, 3, 3]
    """
    rd = p.shape[-1]
    batch_shape = p.shape[: -(num_dim + 1)]
    m = m.reshape(*batch_shape, -1, 1)

    def process_input(p):
        p = p.reshape(*batch_shape, -1, rd)
        p = p * m
        cp = p.sum(dim=-2, keepdim=True) / (
            eps + m.sum(dim=-2, keepdim=True)
        )  # [*, 1, 3]
        p = p - cp
        p = p * m
        return p, cp

    p_rc, cp = process_input(p)  # rc for remove center
    q_rc, cq = process_input(q)
    c = p_rc.transpose(-1, -2) @ q_rc  # [*, 3, 3]
    u, s, v = torch.svd(c)  # [*, 3, 3]
    d = (torch.det(u) * torch.det(v) >= 0.0).type(
        u.dtype
    ) * 2.0 - 1.0  # [*]
    u[..., -1] = u[..., -1] * d[..., None]  # [*, 3]
    u = u @ torch.diag_embed(torch.ones_like(s)) @ v.transpose(-1,-2) # [*, 3, 3]
    u = u.reshape(*batch_shape, *((1,) * num_dim), rd, rd)
    cp = cp.reshape(*batch_shape, *((1,) * num_dim), rd)
    cq = cq.reshape(*batch_shape, *((1,) * num_dim), rd)
    x = cq[..., None, :] - cp[..., None, :] @ u
    return u, x.squeeze(-2)
//  [python]
def get_optimal_transform_v2(
    p: torch.Tensor,
    q: torch.Tensor,
    m: torch.Tensor,
    num_dim: int = 1,
    eps: float = 1e-6,
) -> torch.Tensor:
    """
    calculate u such that p @ u ~ q.
    p, q has shape  [*, *dim, 3]
    mask has shape  [*, *dim]
    ret has shape     [*, *dim, 3, 3]
    """
    rd = p.shape[-1]
    batch_shape = p.shape[: -(num_dim + 1)]
    m = m.reshape(*batch_shape, -1, 1)

    def process_input(p):
        p = p.reshape(*batch_shape, -1, rd)
        p = p * m
        cp = p.sum(dim=-2, keepdim=True) / (
            eps + m.sum(dim=-2, keepdim=True)
        )  # [*, 1, 3]
        p = p - cp
        p = p * m
        return p, cp

    p_rc, cp = process_input(p)  # rc for remove center
    q_rc, cq = process_input(q)
    c = p_rc.transpose(-1, -2) @ q_rc  # [*, 3, 3]
    v, _, w = torch.linalg.svd(c)  # [*, 3, 3]
    d = (torch.linalg.det(v) * torch.linalg.det(w) >= 0.0).type(
        v.dtype
    ) * 2.0 - 1.0  # [*]
    v[..., -1] = v[..., -1] * d[..., None]  # [*, 3]
    u = v @ w  # [*, 3, 3]
    u = u.reshape(*batch_shape, *((1,) * num_dim), rd, rd)
    cp = cp.reshape(*batch_shape, *((1,) * num_dim), rd)
    cq = cq.reshape(*batch_shape, *((1,) * num_dim), rd)
    x = cq[..., None, :] - cp[..., None, :] @ u
    return u, x.squeeze(-2)
//  [python]
def svd_linalg_torch(x):
  x2 = torch.from_numpy(numpy.array(x))
  u, s, v = torch.svd(x2, some=False)
  u1 = u.numpy()
  s1 = s.numpy()
  v1 = v.numpy()
  return u1,s1,v1
//  [python]
def svd_linalg_torch(x):
  x2 = torch.from_numpy(numpy.array(x))
  u, s, v = torch.linalg.svd(x2, full_matrices=False, out=None)
  u1 = u.numpy()
  s1 = s.numpy()
  v1 = v.numpy()
  return u1,s1,v1
//  [python]
def align_torch(source, target):
    '''
    Args:
        source : (B, J, C)
        target : (B, J, C)
        vis:     (B, J, 1)
    '''
    
    device = source.device
    assert len(source.shape) == 3

    muX = torch.mean(target, dim=1, keepdims=True)
    muY = torch.mean(source, dim=1, keepdims=True)
    X0 = target - muX
    Y0 = source - muY

    normX = torch.sqrt(torch.sum(X0**2, dim=(1, 2), keepdims=True))
    normY = torch.sqrt(torch.sum(Y0**2, dim=(1, 2), keepdims=True))
    X0 /= normX
    Y0 /= normY
     
    H = torch.matmul(X0.permute(0, 2, 1), Y0) 
    U, s, V = torch.svd(H)
    Vt = V.permute(0, 2, 1)
    V = Vt.permute(0, 2, 1)
    R = torch.matmul(V, U.permute(0, 2, 1))
    return R
    # Avoid improper rotations (reflections), i.e. rotations with det(R) = -1
    sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))
    V[:, :, -1] *= sign_detR
    R = np.matmul(V, U.transpose(0, 2, 1)) # Rotation
    
    return torch.from_numpy(R).to(device)
//  [python]
def align_torch(source, target):
    '''
    Args:
        source : (B, J, C)
        target : (B, J, C)
        vis:     (B, J, 1)
    '''
    
    device = source.device
    assert len(source.shape) == 3

    muX = torch.mean(target, dim=1, keepdims=True)
    muY = torch.mean(source, dim=1, keepdims=True)
    X0 = target - muX
    Y0 = source - muY

    normX = torch.sqrt(torch.sum(X0**2, dim=(1, 2), keepdims=True))
    normY = torch.sqrt(torch.sum(Y0**2, dim=(1, 2), keepdims=True))
    X0 /= normX
    Y0 /= normY
     
    H = torch.matmul(X0.permute(0, 2, 1), Y0) 
    U, s, Vt = torch.linalg.svd(H)
    V = Vt.permute(0, 2, 1)
    R = torch.matmul(V, U.permute(0, 2, 1))
    return R
    # Avoid improper rotations (reflections), i.e. rotations with det(R) = -1
    sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))
    V[:, :, -1] *= sign_detR
    R = np.matmul(V, U.transpose(0, 2, 1)) # Rotation
    
    return torch.from_numpy(R).to(device)
//  [python]
def vis_body_sigmas(self, packet):
        viz_idx                = packet["viz_idx"]
        world_T_body           = SE3(packet["world_T_body"]).matrix().cpu().numpy()
        world_T_body_cov       = packet["world_T_body_cov"]
        for i in range(len(viz_idx)):
            ix = viz_idx[i].item()
            if ix in self.body_sigma_actors:
                self.viz.remove_geometry(self.body_sigma_actors[ix])
                del self.body_sigma_actors[ix]

            body_pose = world_T_body[i]
            body_cov = world_T_body_cov[i]
            translation_cov = body_cov[3:,3:] # last 3x3 block is translation (gtsam convention)
            U, S, V = torch.svd(translation_cov)
            # TODO: O3D does not support ellipsoids yet...
            body_sigma_actor = self.create_sphere_actor(body_pose, radius=S.max().sqrt().item())

            s = S.max().sqrt().item()
            ic(s)
            max_s = 0.3
            s_norm =  s / max_s if s < max_s else 1.0
            color = self.colormap(s_norm)[:3]
            body_sigma_actor = self.create_camera_actor(color, body_pose, scale=0.02)

            self.body_sigma_actors[ix] = body_sigma_actor
            self.viz.add_geometry(body_sigma_actor, reset_bounding_box=True)
//  [python]
def vis_body_sigmas(self, packet):
        viz_idx                = packet["viz_idx"]
        world_T_body           = SE3(packet["world_T_body"]).matrix().cpu().numpy()
        world_T_body_cov       = packet["world_T_body_cov"]
        for i in range(len(viz_idx)):
            ix = viz_idx[i].item()
            if ix in self.body_sigma_actors:
                self.viz.remove_geometry(self.body_sigma_actors[ix])
                del self.body_sigma_actors[ix]

            body_pose = world_T_body[i]
            body_cov = world_T_body_cov[i]
            translation_cov = body_cov[3:,3:] # last 3x3 block is translation (gtsam convention)
            _, S, _ = torch.linalg.svd(translation_cov)
            # TODO: O3D does not support ellipsoids yet...
            body_sigma_actor = self.create_sphere_actor(body_pose, radius=S.max().sqrt().item())

            s = S.max().sqrt().item()
            ic(s)
            max_s = 0.3
            s_norm =  s / max_s if s < max_s else 1.0
            color = self.colormap(s_norm)[:3]
            body_sigma_actor = self.create_camera_actor(color, body_pose, scale=0.02)

            self.body_sigma_actors[ix] = body_sigma_actor
            self.viz.add_geometry(body_sigma_actor, reset_bounding_box=True)
//  [python]
def eraser(self) -> LeaceEraser:
        """Erasure function lazily computed given the current statistics."""
        eye = torch.eye(self.x_dim, device=self.mean_x.device, dtype=self.mean_x.dtype)

        # Compute the whitening and unwhitening matrices
        if self.method == "leace":
            sigma = self.sigma_xx
            L, V = torch.linalg.eigh(sigma)

            # Threshold used by torch.linalg.pinv
            mask = L > (L[-1] * sigma.shape[-1] * torch.finfo(L.dtype).eps)

            # Assuming PSD; account for numerical error
            L.clamp_min_(0.0)

            W = V * torch.where(mask, L.rsqrt(), 0.0) @ V.mH
            W_inv = V * torch.where(mask, L.sqrt(), 0.0) @ V.mH
        else:
            W, W_inv = eye, eye

        u, s, _ = torch.svd(W @ self.sigma_xz)

        # Throw away singular values that are too small
        u *= s > self.svd_tol

        proj_left = W_inv @ u
        proj_right = u.mH @ W

        if self.constrain_cov_trace and self.method == "leace":
            P = eye - proj_left @ proj_right

            # Prevent the covariance trace from increasing
            sigma = self.sigma_xx
            old_trace = torch.trace(sigma)
            new_trace = torch.trace(P @ sigma @ P.mH)

            # If applying the projection matrix increases the variance, this might
            # cause instability, especially when erasure is applied multiple times.
            # We regularize toward the orthogonal projection matrix to avoid this.
            if new_trace.real > old_trace.real:
                Q = eye - u @ u.mH

                # Set up the variables for the quadratic equation
                x = new_trace
                y = 2 * torch.trace(P @ sigma @ Q.mH)
                z = torch.trace(Q @ sigma @ Q.mH)
                w = old_trace

                # Solve for the mixture of P and Q that makes the trace equal to the
                # trace of the original covariance matrix
                discr = torch.sqrt(
                    4 * w * x - 4 * w * y + 4 * w * z - 4 * x * z + y**2
                )
                alpha1 = (-y / 2 + z - discr / 2) / (x - y + z)
                alpha2 = (-y / 2 + z + discr / 2) / (x - y + z)

                # Choose the positive root
                alpha = torch.where(alpha1.real > 0, alpha1, alpha2).clamp(0, 1)
                P = alpha * P + (1 - alpha) * Q

                # TODO: Avoid using SVD here
                u, s, vh = torch.svd(eye - P)
                proj_left = u * s.sqrt()
                proj_right = vh * s.sqrt()

        return LeaceEraser(
            proj_left, proj_right, bias=self.mean_x if self.affine else None
        )
//  [python]
def eraser(self) -> LeaceEraser:
        """Erasure function lazily computed given the current statistics."""
        eye = torch.eye(self.x_dim, device=self.mean_x.device, dtype=self.mean_x.dtype)

        # Compute the whitening and unwhitening matrices
        if self.method == "leace":
            sigma = self.sigma_xx
            L, V = torch.linalg.eigh(sigma)

            # Threshold used by torch.linalg.pinv
            mask = L > (L[-1] * sigma.shape[-1] * torch.finfo(L.dtype).eps)

            # Assuming PSD; account for numerical error
            L.clamp_min_(0.0)

            W = V * torch.where(mask, L.rsqrt(), 0.0) @ V.mH
            W_inv = V * torch.where(mask, L.sqrt(), 0.0) @ V.mH
        else:
            W, W_inv = eye, eye

        u, s, _ = torch.linalg.svd(W @ self.sigma_xz, full_matrices=False)

        # Throw away singular values that are too small
        u *= s > self.svd_tol

        proj_left = W_inv @ u
        proj_right = u.mH @ W

        if self.constrain_cov_trace and self.method == "leace":
            P = eye - proj_left @ proj_right

            # Prevent the covariance trace from increasing
            sigma = self.sigma_xx
            old_trace = torch.trace(sigma)
            new_trace = torch.trace(P @ sigma @ P.mH)

            # If applying the projection matrix increases the variance, this might
            # cause instability, especially when erasure is applied multiple times.
            # We regularize toward the orthogonal projection matrix to avoid this.
            if new_trace.real > old_trace.real:
                Q = eye - u @ u.mH

                # Set up the variables for the quadratic equation
                x = new_trace
                y = 2 * torch.trace(P @ sigma @ Q.mH)
                z = torch.trace(Q @ sigma @ Q.mH)
                w = old_trace

                # Solve for the mixture of P and Q that makes the trace equal to the
                # trace of the original covariance matrix
                discr = torch.sqrt(
                    4 * w * x - 4 * w * y + 4 * w * z - 4 * x * z + y**2
                )
                alpha1 = (-y / 2 + z - discr / 2) / (x - y + z)
                alpha2 = (-y / 2 + z + discr / 2) / (x - y + z)

                # Choose the positive root
                alpha = torch.where(alpha1.real > 0, alpha1, alpha2).clamp(0, 1)
                P = alpha * P + (1 - alpha) * Q

                # TODO: Avoid using SVD here
                u, s, vh = torch.linalg.svd(eye - P)
                proj_left = u * s.sqrt()
                proj_right = vh * s.sqrt()

        return LeaceEraser(
            proj_left, proj_right, bias=self.mean_x if self.affine else None
        )
//  [python]
def triangulate_simple(points, camera_mats, conf):
    num_cams = len(camera_mats)
    A = torch.zeros((num_cams * 2, 4)).to(points.device)
    for i in range(num_cams):
        x, y = points[i]
        mat = camera_mats[i]
        A[(i * 2):(i * 2 + 1)] = x * mat[2] - mat[0]
        A[(i * 2 + 1):(i * 2 + 2)] = y * mat[2] - mat[1]

    u, s, vh = torch.svd(A)
    p3d = vh[-1]
    p3d = p3d[:3] / p3d[3]
    return p3d
//  [python]
def triangulate_simple(points, camera_mats, conf):
    num_cams = len(camera_mats)
    A = torch.zeros((num_cams * 2, 4)).to(points.device)
    for i in range(num_cams):
        x, y = points[i]
        mat = camera_mats[i]
        A[(i * 2):(i * 2 + 1)] = x * mat[2] - mat[0]
        A[(i * 2 + 1):(i * 2 + 2)] = y * mat[2] - mat[1]

    u, s, vh = torch.linalg.svd(A, full_matrices=True)
    p3d = vh[-1]
    p3d = p3d[:3] / p3d[3]
    return p3d
//  [python]
def forward(
        self,
        inputs: dict,
        training_object: str = "generator",
        training: bool = True,
    ) -> dict:
        X, missing_mask = inputs["X"], inputs["missing_mask"]
        imputation_latent, discrimination, reconstruction, fcn_latent = self.backbone(
            X, missing_mask
        )
        results = {
            "imputation_latent": imputation_latent,
            "discrimination": discrimination,
            "reconstruction": reconstruction,
            "fcn_latent": fcn_latent,
        }

        # return results directly, skip loss calculation to reduce inference time
        if not training:
            return results

        if training_object == "discriminator":
            l_D = F.binary_cross_entropy_with_logits(discrimination, missing_mask)
            results["discrimination_loss"] = l_D
        else:
            # discrimination = discrimination.detach()
            l_G = F.binary_cross_entropy_with_logits(
                discrimination, 1 - missing_mask, weight=1 - missing_mask
            )
            l_pre = calc_mse(imputation_latent, X, missing_mask)
            l_rec = calc_mse(reconstruction, X, missing_mask)
            HTH = torch.matmul(fcn_latent, fcn_latent.permute(1, 0))

            if (
                self.counter_for_updating_F == 0
                or self.counter_for_updating_F % 10 == 0
            ):
                U, s, V = torch.svd(fcn_latent)
                self.term_F = U[:, : self.n_clusters]

            FTHTHF = torch.matmul(
                torch.matmul(self.term_F.permute(1, 0), HTH), self.term_F
            )
            l_kmeans = torch.trace(HTH) - torch.trace(FTHTHF)  # k-means loss
            loss_gene = l_G + l_pre + l_rec + l_kmeans * self.lambda_kmeans
            results["generation_loss"] = loss_gene

        return results
//  [python]
def forward(
        self,
        inputs: dict,
        training_object: str = "generator",
        training: bool = True,
    ) -> dict:
        X, missing_mask = inputs["X"], inputs["missing_mask"]
        imputation_latent, discrimination, reconstruction, fcn_latent = self.backbone(
            X, missing_mask
        )
        results = {
            "imputation_latent": imputation_latent,
            "discrimination": discrimination,
            "reconstruction": reconstruction,
            "fcn_latent": fcn_latent,
        }

        # return results directly, skip loss calculation to reduce inference time
        if not training:
            return results

        if training_object == "discriminator":
            l_D = F.binary_cross_entropy_with_logits(discrimination, missing_mask)
            results["discrimination_loss"] = l_D
        else:
            # discrimination = discrimination.detach()
            l_G = F.binary_cross_entropy_with_logits(
                discrimination, 1 - missing_mask, weight=1 - missing_mask
            )
            l_pre = calc_mse(imputation_latent, X, missing_mask)
            l_rec = calc_mse(reconstruction, X, missing_mask)
            HTH = torch.matmul(fcn_latent, fcn_latent.permute(1, 0))

            if (
                self.counter_for_updating_F == 0
                or self.counter_for_updating_F % 10 == 0
            ):
                U, s, V = torch.linalg.svd(fcn_latent)
                self.term_F = U[:, : self.n_clusters]

            FTHTHF = torch.matmul(
                torch.matmul(self.term_F.permute(1, 0), HTH), self.term_F
            )
            l_kmeans = torch.trace(HTH) - torch.trace(FTHTHF)  # k-means loss
            loss_gene = l_G + l_pre + l_rec + l_kmeans * self.lambda_kmeans
            results["generation_loss"] = loss_gene

        return results
//  [python]
def triangulate_batch(points, camera_mats, scores=None):
    """Given a CxNx2 array of points and Cx4x4 array of extrinsics,
    this returns a Nx3 array of points,
    where N is the number of points and C is the number of cameras.
    Optionally also takes a CxN array of scores for points."""
    # we build an A matrix that is (num_cams, num_points, 2d, 4)
    num_cams = len(camera_mats)

    A_base = camera_mats[:, None, 2:3].broadcast_to(num_cams, 1, 2, 4)
    A = points[:, :, :,None] * A_base - camera_mats[:, None, :2]
    if scores is not None:
        A = A * scores[:, :, None, None]
    # now shape A matrix to (num_points, num_cams*2, 4)
    A = A.swapaxes(0, 1).reshape(-1, num_cams*2, 4)
    u, s, vh = torch.svd(A)
    p3d = vh[:, -1]
    p3d = p3d[:, :3] / p3d[:,3:4]
    return p3d
//  [python]
def triangulate_batch(points, camera_mats, scores=None):
    """Given a CxNx2 array of points and Cx4x4 array of extrinsics,
    this returns a Nx3 array of points,
    where N is the number of points and C is the number of cameras.
    Optionally also takes a CxN array of scores for points."""
    # we build an A matrix that is (num_cams, num_points, 2d, 4)
    num_cams = len(camera_mats)

    A_base = camera_mats[:, None, 2:3].broadcast_to(num_cams, 1, 2, 4)
    A = points[:, :, :,None] * A_base - camera_mats[:, None, :2]
    if scores is not None:
        A = A * scores[:, :, None, None]
    # now shape A matrix to (num_points, num_cams*2, 4)
    A = A.swapaxes(0, 1).reshape(-1, num_cams*2, 4)
    u, s, vh = torch.linalg.svd(A, full_matrices=True)
    p3d = vh[:, -1]
    p3d = p3d[:, :3] / p3d[:,3:4]
    return p3d
//  [python]
def svd(a: ArrayLike, full_matrices=True, compute_uv=True, hermitian=False):
    a = _atleast_float_1(a)
    if not compute_uv:
        return torch.linalg.svdvals(a)

    # NB: ignore the hermitian= argument (no pytorch equivalent)
    result = torch.svd(a, full_matrices=full_matrices)
    return result
//  [python]
def svd(a: ArrayLike, full_matrices=True, compute_uv=True, hermitian=False):
    a = _atleast_float_1(a)
    if not compute_uv:
        return torch.linalg.svdvals(a)

    # NB: ignore the hermitian= argument (no pytorch equivalent)
    result = torch.linalg.svd(a, full_matrices=full_matrices)
    return result
//  [python]
def __init__(self, num_embeddings, embedding_dim, padding_idx=None):
        super(SparseEmbedding, self).__init__(num_embeddings, embedding_dim, padding_idx=padding_idx)
        self.embedding_dim_origin = embedding_dim
        self.embedding_dim_sparsified = embedding_dim
        # 128 is quite good.
        # The k of svd is min(vocab_size, hidden_size)=hidden_size here.
        u, s, vh = torch.svd(self.weight.data.clone().detach())
        self.u = nn.Parameter(u)
        self.s = nn.Parameter(s)
        self.vh = nn.Parameter(vh)
//  [python]
def __init__(self, num_embeddings, embedding_dim, padding_idx=None):
        super(SparseEmbedding, self).__init__(num_embeddings, embedding_dim, padding_idx=padding_idx)
        self.embedding_dim_origin = embedding_dim
        self.embedding_dim_sparsified = embedding_dim
        # 128 is quite good.
        # The k of svd is min(vocab_size, hidden_size)=hidden_size here.
        u, s, vh = torch.linalg.svd(self.weight.data.clone().detach(), full_matrices=False)
        self.u = nn.Parameter(u)
        self.s = nn.Parameter(s)
        self.vh = nn.Parameter(vh)
//  [python]
def get_flattening_scores(self, G, mode='condition_number'):
        if mode == 'condition_number':
            S = torch.svd(G).S
            scores = S.max(1).values/S.min(1).values
        elif mode == 'variance':
            G_mean = torch.mean(G, dim=0, keepdim=True)
            A = torch.inverse(G_mean)@G
            scores = torch.sum(torch.log(torch.svd(A).S)**2, dim=1)
        else:
            pass
        return scores
//  [python]
def get_flattening_scores(self, G, mode='condition_number'):
        if mode == 'condition_number':
            try:
                S = torch.linalg.svd(G).S
            except:
                S = torch.svd(G).S
            scores = S.max(1).values/S.min(1).values
        elif mode == 'variance':
            G_mean = torch.mean(G, dim=0, keepdim=True)
            A = torch.inverse(G_mean)@G
            try:
                scores = torch.sum(torch.log(torch.linalg.svd(A).S)**2, dim=1)
            except:
                scores = torch.sum(torch.log(torch.svd(A).S)**2, dim=1)
        else:
            pass
        return scores
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.svd(0.9 * f_cov + 0.1 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S+1e-5))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach() # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.linalg.svd(0.9 * f_cov + 0.1 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S+1e-5))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach() # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def orthogonal_gradient(self):
        for layer in self.modules():
            if isinstance(layer, SubspaceLayer):
                weight_grad = layer.U.grad
                if not (torch.isnan(weight_grad).any() or torch.isinf(weight_grad).any()):
                    u,s,vh = torch.svd(weight_grad)
                    s[:]=1
                    layer.U.grad =u.mm(torch.diag(s)).mm(vh)
//  [python]
def orthogonal_gradient(self):
        for layer in self.modules():
            if isinstance(layer, SubspaceLayer):
                weight_grad = layer.U.grad
                if not (torch.isnan(weight_grad).any() or torch.isinf(weight_grad).any()):
                    u,s,vh = torch.linalg.svd(weight_grad,full_matrices=False)
                    s[:]=1
                    layer.U.grad =u.mm(torch.diag(s)).mm(vh)
//  [python]
def svd_basis(model, **kwargs):
    """
    Sets the vectors of each slice type to an orthonormal basis.

    :param model: SliceTCA model
    :param kwargs: ignored
    :return: model with new components.
    """

    device = model.device
    ranks = model.ranks

    new_components = [[None, None] for i in range(len(ranks))]
    for i in range(len(ranks)):
        if ranks[i] != 0:
            constructed = model.construct_single_partition(i)
            flattened_constructed = constructed.permute([i]+[q for q in range(len(ranks)) if q != i])
            flattened_constructed = flattened_constructed.reshape(model.dimensions[i],-1).transpose(0,1)

            U, S, V = torch.svd(flattened_constructed.detach().cpu())
            U, S, V = U[:,:ranks[i]], S[:ranks[i]], V[:ranks[i],:]
            U, S, V = U.to(device), S.to(device), V.to(device)

            US = (U @ torch.diag(S))
            slice = US.transpose(0,1).reshape([ranks[i]]+[model.dimensions[q] for q in range(len(ranks)) if q != i])

            new_components[i][0] = V
            new_components[i][1] = slice
        else:
            new_components[i][0] = torch.zeros_like(model.vectors[i][0])
            new_components[i][1] = torch.zeros_like(model.vectors[i][1])

    model.set_components(new_components)

    return model
//  [python]
def svd_basis(model, **kwargs):
    """
    Sets the vectors of each slice type to an orthonormal basis.

    :param model: SliceTCA model
    :param kwargs: ignored
    :return: model with new components.
    """

    device = model.device
    ranks = model.ranks

    new_components = [[None, None] for i in range(len(ranks))]
    for i in range(len(ranks)):
        if ranks[i] != 0:
            constructed = model.construct_single_partition(i)
            flattened_constructed = constructed.permute([i]+[q for q in range(len(ranks)) if q != i])
            flattened_constructed = flattened_constructed.reshape(model.dimensions[i],-1).transpose(0,1)

            U, S, V = torch.linalg.svd(flattened_constructed.detach().cpu(), full_matrices=False)
            U, S, V = U[:,:ranks[i]], S[:ranks[i]], V[:ranks[i]]
            U, S, V = U.to(device), S.to(device), V.to(device)

            US = (U @ torch.diag(S))
            slice = US.transpose(0,1).reshape([ranks[i]]+[model.dimensions[q] for q in range(len(ranks)) if q != i])

            new_components[i][0] = V
            new_components[i][1] = slice
        else:
            new_components[i][0] = torch.zeros_like(model.vectors[i][0])
            new_components[i][1] = torch.zeros_like(model.vectors[i][1])

    model.set_components(new_components)

    return model
//  [python]
def kabsch(x: th.Tensor, y: th.Tensor, weight: Optional[th.Tensor] = None) -> th.Tensor:
    """ Aligns x onto y. x, y are ((...), N, 3) tensors. Weights is ((...), N)
    If rotation fails, at least bring to X to Y's COM
    """
    if weight is None:
        weight = th.ones_like(x[..., 0])

    weight = weight / weight.sum(dim=-1, keepdim=True)
    x_mean = (x * weight[..., None]).sum(dim=-2, keepdim=True)
    y_mean = (y * weight[..., None]).sum(dim=-2, keepdim=True)
    x = x - x_mean
    y = y - y_mean

    # if rotation fails (SVD might fail if matrix is ill-behaved), just bring to same COM
    try:
        # ((...), N, 3) -> ((...), 1, 3, 3)
        cov = th.einsum("...ni,...nj->...ij", x, y * weight[..., None])[..., None, :, :]
        u, s, v = th.svd(cov)
        # Flip the sign of bottom row of each matrix if det product < 0
        det = th.det(v) * th.det(u)
        u_flip = th.ones_like(u)
        u_flip[det < 0, :, -1] = -1.0
        u = u * u_flip
        rot = u @ v
        # align to rotation
        x = rot @ x
    except:
        pass
    return x + y_mean
//  [python]
def kabsch(x: th.Tensor, y: th.Tensor, weight: Optional[th.Tensor] = None) -> th.Tensor:
    """ Aligns x onto y. x, y are ((...), N, 3) tensors. Weights is ((...), N)
    If rotation fails, at least bring to X to Y's COM
    """
    if weight is None:
        weight = th.ones_like(x[..., 0])

    weight = weight / weight.sum(dim=-1, keepdim=True)
    x_mean = (x * weight[..., None]).sum(dim=-2, keepdim=True)
    y_mean = (y * weight[..., None]).sum(dim=-2, keepdim=True)
    x = x - x_mean
    y = y - y_mean

    # if rotation fails (SVD might fail if matrix is ill-behaved), just bring to same COM
    try:
        # ((...), N, 3) -> ((...), 1, 3, 3)
        cov = th.einsum("...ni,...nj->...ij", x, y * weight[..., None])[..., None, :, :]
        u, s, v = th.linalg.svd(cov)
        # Flip the sign of bottom row of each matrix if det product < 0
        det = th.det(v) * th.det(u)
        u_flip = th.ones_like(u)
        u_flip[det < 0, :, -1] = -1.0
        u = u * u_flip
        rot = u @ v
        # align to rotation
        x = rot @ x
    except:
        pass
    return x + y_mean