jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, data):
        """
        Args:
            data: (batch_size, num_samples, num_features) or (num_samples, num_features)
        Returns:
            output: (batch_size, num_samples, num_features) or (num_samples, num_features) in PCA domain.
        """
        n_dims = data.dim()

        assert n_dims in [2, 3], "data is expected 2D or 3D tensor."

        if n_dims == 2:
            data = data.unsqueeze(dim=0) # (batch_size, num_samples, num_features), where batch_size = 1.

        if self.training:
            if self.standardize:
                self.mean, self.std = torch.mean(data, dim=1), torch.std(data, dim=1)
                standardized = self.preprocess(data)
            else:
                standardized = data

            cov = torch.bmm(standardized.permute(0, 2, 1), standardized) / standardized.size(1)
            _, proj_matrix = torch.linalg.eigh(cov) # Computed by assending order
            self.proj_matrix = torch.flip(proj_matrix, dims=(-1,))
        else:
            if self.proj_matrix is None:
                raise RuntimeError("proj_matrix is computed in advance.")

            if self.standardize:
                standardized = self.preprocess(data)
            else:
                standardized = data

        output = torch.bmm(standardized, self.proj_matrix)

        if n_dims == 2:
            output = output.squeeze(dim=0)

        return output
//  [python]
def sym_eig_trunc1(x,nheavyatom,nH,nocc, eig_only=False):

    dtype =  x.dtype
    device = x.device
    if x.dim()==2:
        # e0,v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e0, v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e = torch.zeros((x.shape[0]),dtype=dtype,device=device)
        e[:(nheavyatom*4+nH)] = e0
    else:#need to add large diagonal values to replace 0 padding
        #Gershgorin circle theorem estimate upper bounds of eigenvalues

        # e0, v0 = list(zip(*list(map(
        #                 lambda a,b,c: torch.symeig(pack(a,b,c),eigenvectors=True,upper=True),
        #                 x,nheavyatom, nH))))
        e0, v0 = list(zip(*list(map(
                        lambda a,b,c: torch.symeig(pack(a,b,c),eigenvectors=True,upper=True),
                        x,nheavyatom, nH))))
        if CHECK_DEGENERACY:
            P0 = list(map(
                     lambda e, v, nc : construct_P(e,v,nc),
                     e0, v0,nocc))
        else:
            P0 = list(map(
                     lambda v, nc : 2.0*torch.matmul(v[:,:nc], v[:,:nc].transpose(0,1)),
                     v0,nocc))
        #
        nmol = x.shape[0]
        norb = nheavyatom*4+nH
        e=torch.zeros(x.shape[:2], dtype=dtype, device=device)
        P = torch.zeros_like(x)
        for i in range(nmol):
            e[i,:norb[i]] = e0[i]
            P[i] = unpack(P0[i], nheavyatom[i], nH[i], x.shape[-1])

    if eig_only:
        return e, v0

    # each column of v is a eigenvectors
    # P_alpha_beta = 2.0 * |sum_i c_{i,alpha}*c_{i,beta}, i \in occupied MO

    return e,P, v0
//  [python]
def sym_eig_trunc1(x,nheavyatom,nH,nocc, eig_only=False):

    dtype =  x.dtype
    device = x.device
    if x.dim()==2:
        # e0,v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e0, v = torch.linalg.eigh(pack(x, nheavyatom, nH), UPLO='U')
        e = torch.zeros((x.shape[0]),dtype=dtype,device=device)
        e[:(nheavyatom*4+nH)] = e0
    else:#need to add large diagonal values to replace 0 padding
        #Gershgorin circle theorem estimate upper bounds of eigenvalues

        # e0, v0 = list(zip(*list(map(
        #                 lambda a,b,c: torch.symeig(pack(a,b,c),eigenvectors=True,upper=True),
        #                 x,nheavyatom, nH))))
        e0, v0 = list(zip(*list(map(
                        lambda a,b,c: torch.linalg.eigh(pack(a,b,c), UPLO='U'),
                        x,nheavyatom, nH))))
        if CHECK_DEGENERACY:
            P0 = list(map(
                     lambda e, v, nc : construct_P(e,v,nc),
                     e0, v0,nocc))
        else:
            P0 = list(map(
                     lambda v, nc : 2.0*torch.matmul(v[:,:nc], v[:,:nc].transpose(0,1)),
                     v0,nocc))
        #
        nmol = x.shape[0]
        norb = nheavyatom*4+nH
        e=torch.zeros(x.shape[:2], dtype=dtype, device=device)
        P = torch.zeros_like(x)
        for i in range(nmol):
            e[i,:norb[i]] = e0[i]
            P[i] = unpack(P0[i], nheavyatom[i], nH[i], x.shape[-1])

    if eig_only:
        return e, v0

    # each column of v is a eigenvectors
    # P_alpha_beta = 2.0 * |sum_i c_{i,alpha}*c_{i,beta}, i \in occupied MO

    return e,P, v0
//  [python]
def diagonalize(N, edges=[], antiedges=[], a=1, b=0.3, lamb=0., ptr=None):
        J = torch.zeros((N, N), device=edges.device)  # temporary fix
        for i, j in edges:
            J[i, i] += a
            J[j, j] += a
            J[i, j] = J[j, i] = -a
        for i, j in antiedges:
            J[i, i] -= b
            J[j, j] -= b
            J[i, j] = J[j, i] = b
        J += torch.diag(lamb)
        if ptr is None:
            return torch.symeig(J)

        Ds, Ps = [], []
        for start, end in zip(ptr[:-1], ptr[1:]):
            D, P = torch.symeig(J[start:end, start:end])
            Ds.append(D)
            Ps.append(P)
        return torch.cat(Ds), torch.block_diag(*Ps)
//  [python]
def diagonalize(N, edges=[], antiedges=[], a=1, b=0.3, lamb=0., ptr=None):
        J = torch.zeros((N, N), device=edges.device)  # temporary fix
        for i, j in edges:
            J[i, i] += a
            J[j, j] += a
            J[i, j] = J[j, i] = -a
        for i, j in antiedges:
            J[i, i] -= b
            J[j, j] -= b
            J[i, j] = J[j, i] = b
        J += torch.diag(lamb)
        if ptr is None:
            return torch.linalg.eigh(J)

        Ds, Ps = [], []
        for start, end in zip(ptr[:-1], ptr[1:]):
            D, P = torch.linalg.eigh(J[start:end, start:end])
            Ds.append(D)
            Ps.append(P)
        return torch.cat(Ds), torch.block_diag(*Ps)
//  [python]
def orth_correction(R):
    R[0] /= torch.norm(R[0])
    for i in range(1, R.shape[0]):

        R[i] -= torch.sum( R[:i].t() * torch.matmul(R[:i], R[i]), dim=1)
        R[i] /= torch.norm(R[i])
//  [python]
def orth_correction(R):
    R[0] /= torch.linalg.norm(R[0])
    for i in range(1, R.shape[0]):

        R[i] -= torch.sum( R[:i].t() * torch.matmul(R[:i], R[i]), dim=1)
        R[i] /= torch.linalg.norm(R[i])
//  [python]
def th_batch_kernel_fn(x, y, signal_variance, num_grid_per_dim, raw_length_scale):
    x = x.view(-1, 1, x.shape[-1])
    y = y.view(1, -1, x.shape[-1])
    dist = torch.norm(x - y, dim=2) / (num_grid_per_dim * raw_length_scale)
    return signal_variance * torch.exp(-dist)
//  [python]
def th_batch_kernel_fn(x, y, signal_variance, num_grid_per_dim, raw_length_scale):
    x = x.view(-1, 1, x.shape[-1])
    y = y.view(1, -1, x.shape[-1])
    dist = torch.linalg.norm(x - y, dim=2) / (num_grid_per_dim * raw_length_scale)
    return signal_variance * torch.exp(-dist)
//  [python]
def project_to_plane(points, a, d):
    a = a.reshape((3, 1))
    a = a / torch.norm(a, 2)
    # Project on the same plane but passing through origin
    projections = points - ((points @ a).permute(1, 0) * a).permute(1, 0)

    # shift the points on the plane back to the original d distance
    # from origin
    projections = projections + a.transpose(1, 0) * d
    return projections
//  [python]
def project_to_plane(points, a, d):
    a = a.reshape((3, 1))
    a = a / torch.linalg.norm(a, 2)
    # Project on the same plane but passing through origin
    projections = points - ((points @ a).permute(1, 0) * a).permute(1, 0)

    # shift the points on the plane back to the original d distance
    # from origin
    projections = projections + a.transpose(1, 0) * d
    return projections
//  [python]
def complex_norm(
    c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False
) -> torch.Tensor:
    if not is_complex(c):
        raise TypeError("Input is not a complex tensor.")
    if is_torch_complex_tensor(c):
        return torch.norm(c, dim=dim, keepdim=keepdim)
    else:
        if dim is None:
            return torch.sqrt((c.real**2 + c.imag**2).sum() + EPS)
        else:
            return torch.sqrt(
                (c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS
            )
//  [python]
def complex_norm(
    c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False
) -> torch.Tensor:
    if not is_complex(c):
        raise TypeError("Input is not a complex tensor.")
    if is_torch_complex_tensor(c):
        return torch.linalg.norm(c, dim=dim, keepdim=keepdim)
    else:
        if dim is None:
            return torch.sqrt((c.real**2 + c.imag**2).sum() + EPS)
        else:
            return torch.sqrt(
                (c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS
            )
//  [python]
def pkd_loss(state_S, state_T, mask=None):
    '''
    * Computes normalized vector mse loss at position 0 along `length` dimension. This is the loss used in BERT-PKD, see `Patient Knowledge Distillation for BERT Model Compression <https://arxiv.org/abs/1908.09355>`_.
    * If the hidden sizes of student and teacher are different, 'proj' option is required in `inetermediate_matches` to match the dimensions.

    :param torch.Tensor state_S: tensor of shape  (*batch_size*, *length*, *hidden_size*)
    :param torch.Tensor state_T: tensor of shape  (*batch_size*, *length*, *hidden_size*)
    :param mask: not used.
    '''

    cls_T = state_T[:, 0]  # (batch_size, hidden_dim)
    cls_S = state_S[:, 0]  # (batch_size, hidden_dim)
    normed_cls_T = cls_T / torch.norm(cls_T, dim=1, keepdim=True)
    normed_cls_S = cls_S / torch.norm(cls_S, dim=1, keepdim=True)
    loss = (normed_cls_S - normed_cls_T).pow(2).sum(dim=-1).mean()
    return loss
//  [python]
def pkd_loss(state_S, state_T, mask=None):
    '''
    * Computes normalized vector mse loss at position 0 along `length` dimension. This is the loss used in BERT-PKD, see `Patient Knowledge Distillation for BERT Model Compression <https://arxiv.org/abs/1908.09355>`_.
    * If the hidden sizes of student and teacher are different, 'proj' option is required in `inetermediate_matches` to match the dimensions.

    :param torch.Tensor state_S: tensor of shape  (*batch_size*, *length*, *hidden_size*)
    :param torch.Tensor state_T: tensor of shape  (*batch_size*, *length*, *hidden_size*)
    :param mask: not used.
    '''

    cls_T = state_T[:, 0]  # (batch_size, hidden_dim)
    cls_S = state_S[:, 0]  # (batch_size, hidden_dim)
    normed_cls_T = cls_T / torch.linalg.norm(cls_T, dim=1, keepdim=True)
    normed_cls_S = cls_S / torch.linalg.norm(cls_S, dim=1, keepdim=True)
    loss = (normed_cls_S - normed_cls_T).pow(2).sum(dim=-1).mean()
    return loss
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.linalg.norm(x[l2_cls == c] - tmp_center[c], ord=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.linalg.norm((center_old - center), ord=2, dim=1).max()

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def rand_rotate(dir, ref, pos, alpha=None, device=None):
    if device is None:
        device = 'cpu'
    dir = dir/torch.norm(dir)
    if alpha is None:
        alpha = torch.randn(1).to(device)
    n_pos = pos.shape[0]
    sin, cos = torch.sin(alpha).to(device), torch.cos(alpha).to(device)
    K = 1 - cos
    M = torch.dot(dir, ref)
    nx, ny, nz = dir[0], dir[1], dir[2]
    x0, y0, z0 = ref[0], ref[1], ref[2]
    T = torch.tensor([nx ** 2 * K + cos, nx * ny * K - nz * sin, nx * nz * K + ny * sin,
         (x0 - nx * M) * K + (nz * y0 - ny * z0) * sin,
         nx * ny * K + nz * sin, ny ** 2 * K + cos, ny * nz * K - nx * sin,
         (y0 - ny * M) * K + (nx * z0 - nz * x0) * sin,
         nx * nz * K - ny * sin, ny * nz * K + nx * sin, nz ** 2 * K + cos,
         (z0 - nz * M) * K + (ny * x0 - nx * y0) * sin,
         0, 0, 0, 1], device=device).reshape(4, 4)
    pos = torch.cat([pos.t(), torch.ones(n_pos, device=device).unsqueeze(0)], dim=0)
    rotated_pos = torch.mm(T, pos)[:3]
    return rotated_pos.t()
//  [python]
def rand_rotate(dir, ref, pos, alpha=None, device=None):
    if device is None:
        device = 'cpu'
    dir = dir/torch.linalg.norm(dir)
    if alpha is None:
        alpha = torch.randn(1).to(device)
    n_pos = pos.shape[0]
    sin, cos = torch.sin(alpha).to(device), torch.cos(alpha).to(device)
    K = 1 - cos
    M = torch.dot(dir, ref)
    nx, ny, nz = dir[0], dir[1], dir[2]
    x0, y0, z0 = ref[0], ref[1], ref[2]
    T = torch.tensor([nx ** 2 * K + cos, nx * ny * K - nz * sin, nx * nz * K + ny * sin,
         (x0 - nx * M) * K + (nz * y0 - ny * z0) * sin,
         nx * ny * K + nz * sin, ny ** 2 * K + cos, ny * nz * K - nx * sin,
         (y0 - ny * M) * K + (nx * z0 - nz * x0) * sin,
         nx * nz * K - ny * sin, ny * nz * K + nx * sin, nz ** 2 * K + cos,
         (z0 - nz * M) * K + (ny * x0 - nx * y0) * sin,
         0, 0, 0, 1], device=device).reshape(4, 4)
    pos = torch.cat([pos.t(), torch.ones(n_pos, device=device).unsqueeze(0)], dim=0)
    rotated_pos = torch.mm(T, pos)[:3]
    return rotated_pos.t()
//  [python]
def sym_eig_trunc(x,nheavyatom,nH,nocc, eig_only=False):

    dtype =  x.dtype
    device = x.device
    if x.dim()==2:
        # e0,v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e0, v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e = torch.zeros((x.shape[0]),dtype=dtype,device=device)
        e[:(nheavyatom*4+nH)] = e0
    else:#need to add large diagonal values to replace 0 padding
        #Gershgorin circle theorem estimate upper bounds of eigenvalues
        x0 = pack(x, nheavyatom, nH)
        nmol, size, _ = x0.shape

        aii = x0.diagonal(dim1=1,dim2=2)
        ri = torch.sum(torch.abs(x0),dim=2)-torch.abs(aii)
        hN = torch.max(aii+ri,dim=1)[0]
        dE = hN - torch.min(aii-ri,dim=1)[0] #(maximal - minimal) get range

        norb = nheavyatom*4+nH
        pnorb = size - norb
        nn = torch.max(pnorb).item()
        dx = 0.005
        mutipler = torch.arange(1.0+dx, 1.0+nn*dx+dx, dx, dtype=dtype, device=device)[:nn]
        ind = torch.arange(size, dtype=torch.int64, device=device)
        cond = pnorb>0
        for i in range(nmol):
            if cond[i]:
                x0[i,ind[norb[i]:], ind[norb[i]:]] = mutipler[:pnorb[i]]*dE[i]+hN[i]
        try:
            # e0,v = torch.symeig(x0,eigenvectors=True,upper=True)
            e0, v = torch.symeig(x0,eigenvectors=True,upper=True)
        except:
            if torch.isnan(x0).any():
                print(x0)
            #print(x0.detach().data.numpy())
            # e0,v = torch.symeig(x0,eigenvectors=True,upper=True)
            e0, v = torch.symeig(x0,eigenvectors=True,upper=True)
        e = torch.zeros((nmol, x.shape[-1]),dtype=dtype,device=device)
        e[...,:size] = e0
        for i in range(nmol):
            if cond[i]:
                e[i,norb[i]:size] = 0.0

    if eig_only:
        return e, v

    # each column of v is a eigenvectors
    # P_alpha_beta = 2.0 * |sum_i c_{i,alpha}*c_{i,beta}, i \in occupied MO
    if x.dim()==2:
        if CHECK_DEGENERACY:
            t = construct_P(e, v, nocc)
        else:
            t = 2.0*torch.matmul(v[:,:nocc], v[:,:nocc].transpose(0,1))
    else:
        """
        t = torch.zeros_like(v)
        for i in range(v.shape[0]):
            t[i] = torch.matmul(v[i,:,:nocc[i]], v[i, :,:nocc[i]].transpose(0,1))

        t*=2.0
        """
        if CHECK_DEGENERACY:
            t = torch.stack(list(map(lambda a,b,n : construct_P(a, b, n), e, v, nocc)))
        else:
            #list(map(lambda a,n : print('norm', torch.norm(v, dim=0), n), v, nocc))
            #print(torch.norm())
            t = 2.0*torch.stack(list(map(lambda a,n : torch.matmul(a[:,:n], a[:,:n].transpose(0,1)), v, nocc)))
    P = unpack(t, nheavyatom, nH, x.shape[-1])

    return e,P, v
//  [python]
def sym_eig_trunc(x,nheavyatom,nH,nocc, eig_only=False):

    dtype =  x.dtype
    device = x.device
    if x.dim()==2:
        # e0,v = torch.symeig(pack(x, nheavyatom, nH),eigenvectors=True,upper=True)
        e0, v = torch.linalg.eigh(pack(x, nheavyatom, nH), UPLO='U')
        e = torch.zeros((x.shape[0]),dtype=dtype,device=device)
        e[:(nheavyatom*4+nH)] = e0
    else:#need to add large diagonal values to replace 0 padding
        #Gershgorin circle theorem estimate upper bounds of eigenvalues
        x0 = pack(x, nheavyatom, nH)
        nmol, size, _ = x0.shape

        aii = x0.diagonal(dim1=1,dim2=2)
        ri = torch.sum(torch.abs(x0),dim=2)-torch.abs(aii)
        hN = torch.max(aii+ri,dim=1)[0]
        dE = hN - torch.min(aii-ri,dim=1)[0] #(maximal - minimal) get range

        norb = nheavyatom*4+nH
        pnorb = size - norb
        nn = torch.max(pnorb).item()
        dx = 0.005
        mutipler = torch.arange(1.0+dx, 1.0+nn*dx+dx, dx, dtype=dtype, device=device)[:nn]
        ind = torch.arange(size, dtype=torch.int64, device=device)
        cond = pnorb>0
        for i in range(nmol):
            if cond[i]:
                x0[i,ind[norb[i]:], ind[norb[i]:]] = mutipler[:pnorb[i]]*dE[i]+hN[i]
        try:
            # e0,v = torch.symeig(x0,eigenvectors=True,upper=True)
            e0, v = torch.linalg.eigh(x0, UPLO='U')
        except:
            if torch.isnan(x0).any():
                print(x0)
            #print(x0.detach().data.numpy())
            # e0,v = torch.symeig(x0,eigenvectors=True,upper=True)
            e0, v = torch.linalg.eigh(x0, UPLO='U')
        e = torch.zeros((nmol, x.shape[-1]),dtype=dtype,device=device)
        e[...,:size] = e0
        for i in range(nmol):
            if cond[i]:
                e[i,norb[i]:size] = 0.0

    if eig_only:
        return e, v

    # each column of v is a eigenvectors
    # P_alpha_beta = 2.0 * |sum_i c_{i,alpha}*c_{i,beta}, i \in occupied MO
    if x.dim()==2:
        if CHECK_DEGENERACY:
            t = construct_P(e, v, nocc)
        else:
            t = 2.0*torch.matmul(v[:,:nocc], v[:,:nocc].transpose(0,1))
    else:
        """
        t = torch.zeros_like(v)
        for i in range(v.shape[0]):
            t[i] = torch.matmul(v[i,:,:nocc[i]], v[i, :,:nocc[i]].transpose(0,1))

        t*=2.0
        """
        if CHECK_DEGENERACY:
            t = torch.stack(list(map(lambda a,b,n : construct_P(a, b, n), e, v, nocc)))
        else:
            #list(map(lambda a,n : print('norm', torch.norm(v, dim=0), n), v, nocc))
            #print(torch.norm())
            t = 2.0*torch.stack(list(map(lambda a,n : torch.matmul(a[:,:n], a[:,:n].transpose(0,1)), v, nocc)))
    P = unpack(t, nheavyatom, nH, x.shape[-1])

    return e,P, v
//  [python]
def _quaternion_to_rotation_impl(quaternion: torch.Tensor) -> torch.Tensor:
    check_unit_quaternion(quaternion)

    quaternion = quaternion / torch.norm(quaternion, dim=-1, keepdim=True)
    w = quaternion[..., 0]
    x = quaternion[..., 1]
    y = quaternion[..., 2]
    z = quaternion[..., 3]

    q00 = w * w
    q01 = w * x
    q02 = w * y
    q03 = w * z
    q11 = x * x
    q12 = x * y
    q13 = x * z
    q22 = y * y
    q23 = y * z
    q33 = z * z

    size = quaternion.shape[:-1]
    ret = quaternion.new_zeros(*size, 3, 3)
    ret[..., 0, 0] = q00 + q11 - q22 - q33
    ret[..., 0, 1] = 2 * (q12 - q03)
    ret[..., 0, 2] = 2 * (q13 + q02)
    ret[..., 1, 0] = 2 * (q12 + q03)
    ret[..., 1, 1] = q00 - q11 + q22 - q33
    ret[..., 1, 2] = 2 * (q23 - q01)
    ret[..., 2, 0] = 2 * (q13 - q02)
    ret[..., 2, 1] = 2 * (q23 + q01)
    ret[..., 2, 2] = q00 - q11 - q22 + q33
    return ret
//  [python]
def _quaternion_to_rotation_impl(quaternion: torch.Tensor) -> torch.Tensor:
    check_unit_quaternion(quaternion)

    quaternion = quaternion / torch.linalg.norm(quaternion, dim=-1, keepdim=True)
    w = quaternion[..., 0]
    x = quaternion[..., 1]
    y = quaternion[..., 2]
    z = quaternion[..., 3]

    q00 = w * w
    q01 = w * x
    q02 = w * y
    q03 = w * z
    q11 = x * x
    q12 = x * y
    q13 = x * z
    q22 = y * y
    q23 = y * z
    q33 = z * z

    size = quaternion.shape[:-1]
    ret = quaternion.new_zeros(*size, 3, 3)
    ret[..., 0, 0] = q00 + q11 - q22 - q33
    ret[..., 0, 1] = 2 * (q12 - q03)
    ret[..., 0, 2] = 2 * (q13 + q02)
    ret[..., 1, 0] = 2 * (q12 + q03)
    ret[..., 1, 1] = q00 - q11 + q22 - q33
    ret[..., 1, 2] = 2 * (q23 - q01)
    ret[..., 2, 0] = 2 * (q13 - q02)
    ret[..., 2, 1] = 2 * (q23 + q01)
    ret[..., 2, 2] = q00 - q11 - q22 + q33
    return ret
//  [python]
def compute_mpjpe(predicted, target, valid_mask=None, pck_joints=None, sample_wise=True):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape, print(predicted.shape, target.shape)
    mpjpe = torch.norm(predicted - target, p=2, dim=-1)
    
    if pck_joints is None:
        if sample_wise:
            mpjpe_batch = (mpjpe*valid_mask.float()).sum(-1)/valid_mask.float().sum(-1) if valid_mask is not None else mpjpe.mean(-1)
        else:
            mpjpe_batch = mpjpe[valid_mask] if valid_mask is not None else mpjpe
        return mpjpe_batch
    else:
        mpjpe_pck_batch = mpjpe[:,pck_joints]
        return mpjpe_pck_batch
//  [python]
def compute_mpjpe(predicted, target, valid_mask=None, pck_joints=None, sample_wise=True):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape, print(predicted.shape, target.shape)
    mpjpe = torch.linalg.norm(predicted - target, ord=2, dim=-1)
    
    if pck_joints is None:
        if sample_wise:
            mpjpe_batch = (mpjpe*valid_mask.float()).sum(-1)/valid_mask.float().sum(-1) if valid_mask is not None else mpjpe.mean(-1)
        else:
            mpjpe_batch = mpjpe[valid_mask] if valid_mask is not None else mpjpe
        return mpjpe_batch
    else:
        mpjpe_pck_batch = mpjpe[:,pck_joints]
        return mpjpe_pck_batch
//  [python]
def _calc_matched_PCKh_(real, pred, kp2d_mask, error_thresh=0.143):
    # error_thresh is set as the ratio between the head and the body.
    # he head / body for normal people are between 6~8, therefore, we set it to 1/7=0.143
    PCKs = torch.ones(len(kp2d_mask)).float().cuda()*-1.
    if kp2d_mask.sum()>0:
        vis = (real>-1.).sum(-1)==real.shape[-1]
        error = torch.norm(real-pred, p=2, dim=-1)
        
        for ind, (e, v) in enumerate(zip(error, vis)):
            if v.sum() < 2:
                continue
            real_valid = real[ind,v]
            person_scales = torch.sqrt((real_valid[:,0].max(-1).values - real_valid[:,0].min(-1).values)**2 + \
                            (real_valid[:,1].max(-1).values - real_valid[:,1].min(-1).values)**2)
            error_valid = e[v]
            correct_kp_mask = ((error_valid / person_scales) < error_thresh).float()
            PCKs[ind] = correct_kp_mask.sum()/len(correct_kp_mask)
    return PCKs
//  [python]
def _calc_matched_PCKh_(real, pred, kp2d_mask, error_thresh=0.143):
    # error_thresh is set as the ratio between the head and the body.
    # he head / body for normal people are between 6~8, therefore, we set it to 1/7=0.143
    PCKs = torch.ones(len(kp2d_mask)).float().cuda()*-1.
    if kp2d_mask.sum()>0:
        vis = (real>-1.).sum(-1)==real.shape[-1]
        error = torch.linalg.norm(real-pred, ord=2, dim=-1)
        
        for ind, (e, v) in enumerate(zip(error, vis)):
            if v.sum() < 2:
                continue
            real_valid = real[ind,v]
            person_scales = torch.sqrt((real_valid[:,0].max(-1).values - real_valid[:,0].min(-1).values)**2 + \
                            (real_valid[:,1].max(-1).values - real_valid[:,1].min(-1).values)**2)
            error_valid = e[v]
            correct_kp_mask = ((error_valid / person_scales) < error_thresh).float()
            PCKs[ind] = correct_kp_mask.sum()/len(correct_kp_mask)
    return PCKs
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4] 
            - src_keypts:    [bs, num_corr, 3]    
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:    
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 0.1:
            inlier_threshold_list = [0.1] * 20
        else:
            inlier_threshold_list = [1.2] * 20
        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)
            L2_dis = torch.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(A=src_keypts[:, pred_inlier, :], B=tgt_keypts[:, pred_inlier, :], weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier])
        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4] 
            - src_keypts:    [bs, num_corr, 3]    
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:    
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 0.1:
            inlier_threshold_list = [0.1] * 20
        else:
            inlier_threshold_list = [1.2] * 20
        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)
            L2_dis = torch.linalg.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(A=src_keypts[:, pred_inlier, :], B=tgt_keypts[:, pred_inlier, :], weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier])
        return initial_trans
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input: 
            - M:          [bs, num_corr, num_corr] the compatibility matrix 
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output: 
            - confidence  
        """
        if method == 'eig_value':
            max_eig_value = leading_eig[:, None, :] @ M @ leading_eig[:, :, None] / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            max_eig_value = leading_eig[:, None, :] @ M @ leading_eig[:, :, None] / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.norm(solution, dim=1, keepdim=True) + 1e-06)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = second_eig[:, None, :] @ B @ second_eig[:, :, None] / (second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input: 
            - M:          [bs, num_corr, num_corr] the compatibility matrix 
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output: 
            - confidence  
        """
        if method == 'eig_value':
            max_eig_value = leading_eig[:, None, :] @ M @ leading_eig[:, :, None] / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            max_eig_value = leading_eig[:, None, :] @ M @ leading_eig[:, :, None] / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.linalg.norm(solution, dim=1, keepdim=True) + 1e-06)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = second_eig[:, None, :] @ B @ second_eig[:, :, None] / (second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def _jquaternion_to_rotation_impl(
    quaternion: torch.Tensor,
) -> Tuple[List[torch.Tensor], torch.Tensor]:
    if quaternion.ndim == 1:
        quaternion = quaternion.unsqueeze(0)
    check_unit_quaternion(quaternion)

    quaternion_norm = torch.norm(quaternion, dim=-1, keepdim=True)
    quaternion = quaternion / quaternion_norm
    w = quaternion[..., 0]
    x = quaternion[..., 1]
    y = quaternion[..., 2]
    z = quaternion[..., 3]

    q00 = w * w
    q01 = w * x
    q02 = w * y
    q03 = w * z
    q11 = x * x
    q12 = x * y
    q13 = x * z
    q22 = y * y
    q23 = y * z
    q33 = z * z

    size = quaternion.shape[:-1]
    ret = quaternion.new_zeros(*size, 3, 3)
    ret[..., 0, 0] = q00 + q11 - q22 - q33
    ret[..., 0, 1] = 2 * (q12 - q03)
    ret[..., 0, 2] = 2 * (q13 + q02)
    ret[..., 1, 0] = 2 * (q12 + q03)
    ret[..., 1, 1] = q00 - q11 + q22 - q33
    ret[..., 1, 2] = 2 * (q23 - q01)
    ret[..., 2, 0] = 2 * (q13 - q02)
    ret[..., 2, 1] = 2 * (q23 + q01)
    ret[..., 2, 2] = q00 - q11 - q22 + q33

    temp = -2 * quaternion / quaternion_norm
    jac = quaternion.new_zeros(*size, 3, 4)
    jac[..., :, :1] = temp[..., 1:].view(*size, 3, 1)
    jac[..., :, 1:] = _hat_autograd_fn(temp[..., 1:])
    jac[..., 0, 1] = -temp[..., 0]
    jac[..., 1, 2] = -temp[..., 0]
    jac[..., 2, 3] = -temp[..., 0]

    return [jac], ret
//  [python]
def _jquaternion_to_rotation_impl(
    quaternion: torch.Tensor,
) -> Tuple[List[torch.Tensor], torch.Tensor]:
    if quaternion.ndim == 1:
        quaternion = quaternion.unsqueeze(0)
    check_unit_quaternion(quaternion)

    quaternion_norm = torch.linalg.norm(quaternion, dim=-1, keepdim=True)
    quaternion = quaternion / quaternion_norm
    w = quaternion[..., 0]
    x = quaternion[..., 1]
    y = quaternion[..., 2]
    z = quaternion[..., 3]

    q00 = w * w
    q01 = w * x
    q02 = w * y
    q03 = w * z
    q11 = x * x
    q12 = x * y
    q13 = x * z
    q22 = y * y
    q23 = y * z
    q33 = z * z

    size = quaternion.shape[:-1]
    ret = quaternion.new_zeros(*size, 3, 3)
    ret[..., 0, 0] = q00 + q11 - q22 - q33
    ret[..., 0, 1] = 2 * (q12 - q03)
    ret[..., 0, 2] = 2 * (q13 + q02)
    ret[..., 1, 0] = 2 * (q12 + q03)
    ret[..., 1, 1] = q00 - q11 + q22 - q33
    ret[..., 1, 2] = 2 * (q23 - q01)
    ret[..., 2, 0] = 2 * (q13 - q02)
    ret[..., 2, 1] = 2 * (q23 + q01)
    ret[..., 2, 2] = q00 - q11 - q22 + q33

    temp = -2 * quaternion / quaternion_norm
    jac = quaternion.new_zeros(*size, 3, 4)
    jac[..., :, :1] = temp[..., 1:].view(*size, 3, 1)
    jac[..., :, 1:] = _hat_autograd_fn(temp[..., 1:])
    jac[..., 0, 1] = -temp[..., 0]
    jac[..., 1, 2] = -temp[..., 0]
    jac[..., 2, 3] = -temp[..., 0]

    return [jac], ret
//  [python]
def batch_rodrigues(axisang):
    # This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L37
    # axisang N x 3
    axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def batch_rodrigues(axisang):
    # This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L37
    # axisang N x 3
    axisang_norm = torch.linalg.norm(axisang + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def forward(self, trans, gt_trans, src_keypts, tgt_keypts, probs):
        """
        Transformation Loss
        Inputs:
            - trans:      [bs, 4, 4] SE3 transformation matrices
            - gt_trans:   [bs, 4, 4] ground truth SE3 transformation matrices
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
            - probs:     [bs, num_corr] predicted inlier probability
        Outputs:
            - loss     transformation loss 
            - recall   registration recall (re < re_thre & te < te_thre)
            - RE       rotation error 
            - TE       translation error
            - RMSE     RMSE under the predicted transformation
        """
        bs = trans.shape[0]
        R, t = decompose_trans(trans)
        gt_R, gt_t = decompose_trans(gt_trans)
        recall = 0
        RE = torch.tensor(0.0)
        TE = torch.tensor(0.0)
        RMSE = torch.tensor(0.0)
        loss = torch.tensor(0.0)
        for i in range(bs):
            re = torch.acos(torch.clamp((torch.trace(R[i].T @ gt_R[i]) - 1) / 2.0, min=-1, max=1))
            te = torch.sqrt(torch.sum((t[i] - gt_t[i]) ** 2))
            warp_src_keypts = transform(src_keypts[i], trans[i])
            rmse = torch.norm(warp_src_keypts - tgt_keypts, dim=-1).mean()
            re = re * 180 / np.pi
            te = te * 100
            if te < self.te_thre and re < self.re_thre:
                recall += 1
            RE += re
            TE += te
            RMSE += rmse
            pred_inliers = torch.where(probs[i] > 0)[0]
            if len(pred_inliers) < 1:
                loss += torch.tensor(0.0)
            else:
                warp_src_keypts = transform(src_keypts[i], trans[i])
                loss += ((warp_src_keypts - tgt_keypts) ** 2).sum(-1).mean()
        return loss / bs, recall * 100.0 / bs, RE / bs, TE / bs, RMSE / bs
//  [python]
def forward(self, trans, gt_trans, src_keypts, tgt_keypts, probs):
        """
        Transformation Loss
        Inputs:
            - trans:      [bs, 4, 4] SE3 transformation matrices
            - gt_trans:   [bs, 4, 4] ground truth SE3 transformation matrices
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
            - probs:     [bs, num_corr] predicted inlier probability
        Outputs:
            - loss     transformation loss 
            - recall   registration recall (re < re_thre & te < te_thre)
            - RE       rotation error 
            - TE       translation error
            - RMSE     RMSE under the predicted transformation
        """
        bs = trans.shape[0]
        R, t = decompose_trans(trans)
        gt_R, gt_t = decompose_trans(gt_trans)
        recall = 0
        RE = torch.tensor(0.0)
        TE = torch.tensor(0.0)
        RMSE = torch.tensor(0.0)
        loss = torch.tensor(0.0)
        for i in range(bs):
            re = torch.acos(torch.clamp((torch.trace(R[i].T @ gt_R[i]) - 1) / 2.0, min=-1, max=1))
            te = torch.sqrt(torch.sum((t[i] - gt_t[i]) ** 2))
            warp_src_keypts = transform(src_keypts[i], trans[i])
            rmse = torch.linalg.norm(warp_src_keypts - tgt_keypts, dim=-1).mean()
            re = re * 180 / np.pi
            te = te * 100
            if te < self.te_thre and re < self.re_thre:
                recall += 1
            RE += re
            TE += te
            RMSE += rmse
            pred_inliers = torch.where(probs[i] > 0)[0]
            if len(pred_inliers) < 1:
                loss += torch.tensor(0.0)
            else:
                warp_src_keypts = transform(src_keypts[i], trans[i])
                loss += ((warp_src_keypts - tgt_keypts) ** 2).sum(-1).mean()
        return loss / bs, recall * 100.0 / bs, RE / bs, TE / bs, RMSE / bs
//  [python]
def init_meta(self) -> tp.Dict[str, np.ndarray]:
        if self.solved_meta is not None:
            return self.solved_meta
        task = torch.randn(self.sf_dim)
        task = task / torch.norm(task)
        task_array = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task_array
        return meta
//  [python]
def init_meta(self) -> tp.Dict[str, np.ndarray]:
        if self.solved_meta is not None:
            return self.solved_meta
        task = torch.randn(self.sf_dim)
        task = task / torch.linalg.norm(task)
        task_array = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task_array
        return meta
//  [python]
def compute_intr_reward(self, task, next_obs, step) -> Tuple[Any, Any]:
        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
        reward = self.pbe(rep)
        intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", task, rep).reshape(-1, 1)

        return intr_ent_reward, intr_sf_reward
//  [python]
def compute_intr_reward(self, task, next_obs, step) -> Tuple[Any, Any]:
        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
        reward = self.pbe(rep)
        intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.linalg.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", task, rep).reshape(-1, 1)

        return intr_ent_reward, intr_sf_reward
//  [python]
def batch_rodrigues2(axisang):
    axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle + 1e-8)

    rx, ry, rz = torch.split(axisang_normalized, 1, dim=1)
    zeros = torch.zeros_like(rx)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((axisang.shape[0], 3, 3))
    ident = torch.eye(3, dtype=axisang.dtype,
                      device=axisang.device).unsqueeze(dim=0)

    angle = angle[:, :, None]
    rot_mat = ident + torch.sin(angle) * K + \
        (1 - torch.cos(angle)) * torch.bmm(K, K)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def batch_rodrigues2(axisang):
    axisang_norm = torch.linalg.norm(axisang + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle + 1e-8)

    rx, ry, rz = torch.split(axisang_normalized, 1, dim=1)
    zeros = torch.zeros_like(rx)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((axisang.shape[0], 3, 3))
    ident = torch.eye(3, dtype=axisang.dtype,
                      device=axisang.device).unsqueeze(dim=0)

    angle = angle[:, :, None]
    rot_mat = ident + torch.sin(angle) * K + \
        (1 - torch.cos(angle)) * torch.bmm(K, K)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def forward(self, data):
        """
        Input:
            - corr_pos:   [bs, num_corr, 6]
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
            - testing:    flag for test phase, if False will not calculate M and post-refinement.
        Output: (dict)
            - final_trans:   [bs, 4, 4], the predicted transformation matrix. 
            - final_labels:  [bs, num_corr], the predicted inlier/outlier label (0,1), for classification loss calculation.
            - M:             [bs, num_corr, num_corr], feature similarity matrix, for SM loss calculation.
            - seed_trans:    [bs, num_seeds, 4, 4],  the predicted transformation matrix associated with each seeding point, deprecated.
            - corr_features: [bs, num_corr, num_channels], the feature for each correspondence, for circle loss calculation, deprecated.
            - confidence:    [bs], confidence of returned results, for safe guard, deprecated.
        """
        corr_pos, src_keypts, tgt_keypts = data['corr_pos'], data['src_keypts'], data['tgt_keypts']
        bs, num_corr = corr_pos.shape[0], corr_pos.shape[1]
        testing = 'testing' in data.keys()
        with torch.no_grad():
            src_dist = torch.norm(src_keypts[:, :, None, :] - src_keypts[:, None, :, :], dim=-1)
            corr_compatibility = src_dist - torch.norm(tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :], dim=-1)
            corr_compatibility = torch.clamp(1.0 - corr_compatibility ** 2 / self.sigma_spat ** 2, min=0)
        corr_features = self.encoder(corr_pos.permute(0, 2, 1), corr_compatibility).permute(0, 2, 1)
        normed_corr_features = F.normalize(corr_features, p=2, dim=-1)
        if not testing:
            M = torch.matmul(normed_corr_features, normed_corr_features.permute(0, 2, 1))
            M = torch.clamp(1 - (1 - M) / self.sigma ** 2, min=0, max=1)
            M[:, torch.arange(M.shape[1]), torch.arange(M.shape[1])] = 0
        else:
            M = None
        confidence = self.classification(corr_features.permute(0, 2, 1)).squeeze(1)
        if testing:
            seeds = self.pick_seeds(src_dist, confidence, R=self.nms_radius, max_num=int(num_corr * self.ratio))
        else:
            seeds = torch.argsort(confidence, dim=1, descending=True)[:, 0:int(num_corr * self.ratio)]
        seed_trans, seed_fitness, final_trans, final_labels = self.cal_seed_trans(seeds, normed_corr_features, src_keypts, tgt_keypts)
        if testing:
            final_trans = self.post_refinement(final_trans, src_keypts, tgt_keypts)
        if not testing:
            final_labels = confidence
        res = {'final_trans': final_trans, 'final_labels': final_labels, 'M': M}
        return res