jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])

    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda') 
    
    Xf = torch.rand(Nf, 2)*2-1
    

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.linalg.norm(Xf, dim=1)>=delta)] 
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.linalg.norm(Xf, dim=1)>=delta)]  
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2, requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def batch_rodrigues(self, theta):
            l1norm = torch.norm(theta + 1e-8, p=2, dim=1)
            angle = torch.unsqueeze(l1norm, -1)
            normalized = torch.div(theta, angle)
            angle = angle * 0.5
            v_cos = torch.cos(angle)
            v_sin = torch.sin(angle)
            quat = self.quat2mat(torch.cat([v_cos, v_sin * normalized], dim=1))
            return quat
//  [python]
def batch_rodrigues(self, theta):
            l1norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
            angle = torch.unsqueeze(l1norm, -1)
            normalized = torch.div(theta, angle)
            angle = angle * 0.5
            v_cos = torch.cos(angle)
            v_sin = torch.sin(angle)
            quat = self.quat2mat(torch.cat([v_cos, v_sin * normalized], dim=1))
            return quat
//  [python]
def batch_rodrigues(theta):
    l1norm = torch.norm(theta + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = quat2mat(torch.cat([v_cos, v_sin * normalized], dim=1))
    return quat
//  [python]
def batch_rodrigues(theta):
    l1norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = quat2mat(torch.cat([v_cos, v_sin * normalized], dim=1))
    return quat
//  [python]
def fit(self, sparse_thetas: TensorList, time_derivs: TensorList) -> TensorList:
        """Calculates the coefficients of the constraint using the QR decomposition for every pair
        of sparse feature matrix and time derivative.

        Args:
            sparse_thetas (TensorList): List containing the sparse feature tensors.
            time_derivs (TensorList): List containing the time derivatives.

        Returns:
            (TensorList): Calculated coefficients.
        """
        coeff_vectors = []
        for theta, dt in zip(sparse_thetas, time_derivs):
            # We use the augmented data method
            norm = torch.norm(theta, dim=0)
            l_normed = torch.diag(
                torch.sqrt(self.l) * norm
            )  # we norm l rather than theta cause we want unnormed coeffs out
            X = torch.cat((theta, l_normed), dim=0)
            y = torch.cat((dt, torch.zeros((theta.shape[1], 1))), dim=0)
            # Now solve like normal OLS prob lem
            Q, R = torch.qr(X)  # solution of lst. sq. by QR decomp.
            coeff_vectors.append(torch.inverse(R) @ Q.T @ y)

        return coeff_vectors
//  [python]
def fit(self, sparse_thetas: TensorList, time_derivs: TensorList) -> TensorList:
        """Calculates the coefficients of the constraint using the QR decomposition for every pair
        of sparse feature matrix and time derivative.

        Args:
            sparse_thetas (TensorList): List containing the sparse feature tensors.
            time_derivs (TensorList): List containing the time derivatives.

        Returns:
            (TensorList): Calculated coefficients.
        """
        coeff_vectors = []
        for theta, dt in zip(sparse_thetas, time_derivs):
            # We use the augmented data method
            norm = torch.linalg.norm(theta, dim=0)
            l_normed = torch.diag(
                torch.sqrt(self.l) * norm
            )  # we norm l rather than theta cause we want unnormed coeffs out
            X = torch.cat((theta, l_normed), dim=0)
            y = torch.cat((dt, torch.zeros((theta.shape[1], 1))), dim=0)
            # Now solve like normal OLS prob lem
            Q, R = torch.qr(X)  # solution of lst. sq. by QR decomp.
            coeff_vectors.append(torch.inverse(R) @ Q.T @ y)

        return coeff_vectors
//  [python]
def calculate_face_normal(vertices, faces):
    """
    Calculate face normal from vertices and faces.
    Vertices has shape [N, 3], faces has shape [M, 3]
    Output has shape [M, 3].
    """
    face_pos = vertices[[faces]] #shape [M, 3, 3]

    BA = face_pos[:, 1] - face_pos[:, 0]
    BA = BA / torch.norm(BA, dim=-1, keepdim=True) # Prevent too small values
    CA = face_pos[:, 2] - face_pos[:, 0]
    CA = CA / torch.norm(CA, dim=-1, keepdim=True)

    normals = torch.cross(BA, CA, dim=-1)

    normals = normals / torch.norm(normals, dim=-1, keepdim=True)

    return normals
//  [python]
def calculate_face_normal(vertices, faces):
    """
    Calculate face normal from vertices and faces.
    Vertices has shape [N, 3], faces has shape [M, 3]
    Output has shape [M, 3].
    """
    face_pos = vertices[[faces]] #shape [M, 3, 3]

    BA = face_pos[:, 1] - face_pos[:, 0]
    BA = BA / torch.linalg.norm(BA, dim=-1, keepdim=True) # Prevent too small values
    CA = face_pos[:, 2] - face_pos[:, 0]
    CA = CA / torch.linalg.norm(CA, dim=-1, keepdim=True)

    normals = torch.cross(BA, CA, dim=-1)

    normals = normals / torch.linalg.norm(normals, dim=-1, keepdim=True)

    return normals
//  [python]
def _euclid_distance(x, y, dim=-1):
    return torch.norm(x - y, dim=dim)
//  [python]
def _euclid_distance(x, y, dim=-1):
    return torch.linalg.norm(x - y, dim=dim)
//  [python]
def _init_kmeans(self, normalized_data):
        """
        Args:
            normalized_data: (batch_size, num_samples, num_features)
        Returns:
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        """
        cluster_ids, centroids = super()._init_kmeans(normalized_data)
        centroids = centroids / torch.norm(centroids, dim=2, keepdim=True)

        return cluster_ids, centroids
//  [python]
def _init_kmeans(self, normalized_data):
        """
        Args:
            normalized_data: (batch_size, num_samples, num_features)
        Returns:
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        """
        cluster_ids, centroids = super()._init_kmeans(normalized_data)
        centroids = centroids / torch.linalg.norm(centroids, dim=2, keepdim=True)

        return cluster_ids, centroids
//  [python]
def post_refinement(initial_trans, src_kpts, tgt_kpts, iters, weights=None):
    inlier_threshold = 0.1
    pre_inlier_count = 0
    for i in range(iters):
        pred_tgt = transform(src_kpts, initial_trans)
        L2_dis = torch.norm(pred_tgt - tgt_kpts, dim=-1)
        pred_inlier = (L2_dis < inlier_threshold)[0]
        inlier_count = torch.sum(pred_inlier)
        if inlier_count <= pre_inlier_count:
            break
        pre_inlier_count = inlier_count
        initial_trans = rigid_transform_3d(
            A=src_kpts[:, pred_inlier, :],
            B=tgt_kpts[:, pred_inlier, :],
            weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier]
        )
    return initial_trans
//  [python]
def post_refinement(initial_trans, src_kpts, tgt_kpts, iters, weights=None):
    inlier_threshold = 0.1
    pre_inlier_count = 0
    for i in range(iters):
        pred_tgt = transform(src_kpts, initial_trans)
        L2_dis = torch.linalg.norm(pred_tgt - tgt_kpts, dim=-1)
        pred_inlier = (L2_dis < inlier_threshold)[0]
        inlier_count = torch.sum(pred_inlier)
        if inlier_count <= pre_inlier_count:
            break
        pre_inlier_count = inlier_count
        initial_trans = rigid_transform_3d(
            A=src_kpts[:, pred_inlier, :],
            B=tgt_kpts[:, pred_inlier, :],
            weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier]
        )
    return initial_trans
//  [python]
def Replicator(A_2):
    l = A_2.size(0)
    x = torch.ones(l, 1) / l

    x = Variable(x, requires_grad=True).cpu()
    A_f = A_2.cpu()
    x = x.cpu()

    # x = x * (torch.mm(A_f, x))
    # x = x / torch.sum(x)

    toll = 0.0000001
    ero = 2 * toll + 1
    max_it = 5
    if max_it:
        # print(self.max_it)
        max_it = max_it
    else:
        max_it = float('inf')

    count = int(0)

    x_old = x.cpu()
    x_old = x_old.type(torch.FloatTensor)
    x = (x * (torch.matmul(A_f, x))).contiguous()

    xx = torch.norm(x, p=2, dim=0).detach()
    x = x.div(xx.expand_as(x))

    while ero > toll and count < max_it:
        x_old = x.cpu()
        x_old = x_old.type(torch.FloatTensor)
        x = (x * (torch.matmul(A_f, x))).contiguous()

        xx = torch.norm(x, p=2, dim=0).detach()
        x = x.div(xx.expand_as(x))

        ero = torch.norm(x - x_old)
        ero = float(ero)
        count = count + 1

    return x.cuda()
//  [python]
def Replicator(A_2):
    l = A_2.size(0)
    x = torch.ones(l, 1) / l

    x = Variable(x, requires_grad=True).cpu()
    A_f = A_2.cpu()
    x = x.cpu()

    # x = x * (torch.mm(A_f, x))
    # x = x / torch.sum(x)

    toll = 0.0000001
    ero = 2 * toll + 1
    max_it = 5
    if max_it:
        # print(self.max_it)
        max_it = max_it
    else:
        max_it = float('inf')

    count = int(0)

    x_old = x.cpu()
    x_old = x_old.type(torch.FloatTensor)
    x = (x * (torch.matmul(A_f, x))).contiguous()

    xx = torch.linalg.norm(x, ord=2, dim=0).detach()
    x = x.div(xx.expand_as(x))

    while ero > toll and count < max_it:
        x_old = x.cpu()
        x_old = x_old.type(torch.FloatTensor)
        x = (x * (torch.matmul(A_f, x))).contiguous()

        xx = torch.linalg.norm(x, ord=2, dim=0).detach()
        x = x.div(xx.expand_as(x))

        ero = torch.linalg.norm(x - x_old)
        ero = float(ero)
        count = count + 1

    return x.cuda()
//  [python]
def update_once(self, normalized_data, cluster_ids=None, centroids=None):
        """
        Args:
            normalized_data: (batch_size, num_samples, num_features)
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        Returns:
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        """
        K = self.K
        mask = F.one_hot(cluster_ids, num_classes=K) # (batch_size, num_samples, K)

        """
        1. Calculate centroids
        """
        masked_data = mask.unsqueeze(dim=3) * normalized_data.unsqueeze(dim=2) # (batch_size, num_samples, K, num_features)
        pseudo_centroids = masked_data.sum(dim=1) # (batch_size, K, num_features)
        denominator = mask.sum(dim=1).unsqueeze(dim=2) # (batch_size, K, 1)
        centroids = pseudo_centroids / denominator # (batch_size, K, num_features)
        centroids = centroids / torch.norm(centroids, dim=2, keepdim=True)

        """
        2. Put labels based on distance
        """
        distance = _neg_dot_product(normalized_data.unsqueeze(dim=2), centroids.unsqueeze(dim=1), dim=3) # (batch_size, num_samples, K)
        cluster_ids = torch.argmin(distance, dim=2) # (batch_size, num_samples)

        return cluster_ids, centroids
//  [python]
def update_once(self, normalized_data, cluster_ids=None, centroids=None):
        """
        Args:
            normalized_data: (batch_size, num_samples, num_features)
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        Returns:
            cluster_ids: (batch_size, num_samples)
            centroids: (batch_size, K, num_features)
        """
        K = self.K
        mask = F.one_hot(cluster_ids, num_classes=K) # (batch_size, num_samples, K)

        """
        1. Calculate centroids
        """
        masked_data = mask.unsqueeze(dim=3) * normalized_data.unsqueeze(dim=2) # (batch_size, num_samples, K, num_features)
        pseudo_centroids = masked_data.sum(dim=1) # (batch_size, K, num_features)
        denominator = mask.sum(dim=1).unsqueeze(dim=2) # (batch_size, K, 1)
        centroids = pseudo_centroids / denominator # (batch_size, K, num_features)
        centroids = centroids / torch.linalg.norm(centroids, dim=2, keepdim=True)

        """
        2. Put labels based on distance
        """
        distance = _neg_dot_product(normalized_data.unsqueeze(dim=2), centroids.unsqueeze(dim=1), dim=3) # (batch_size, num_samples, K)
        cluster_ids = torch.argmin(distance, dim=2) # (batch_size, num_samples)

        return cluster_ids, centroids
//  [python]
def get_dihedral_torch(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from: 
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Can't use torch.dot bc it does not broadcast
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return torch.atan2( ( (torch.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,  
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )
//  [python]
def get_dihedral_torch(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from: 
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Can't use torch.dot bc it does not broadcast
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return torch.atan2( ( (torch.linalg.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,  
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )
//  [python]
def forward(self, data, iteration=None):
        """
        Args:
            data <torch.Tensor>: (batch_size, num_samples, num_features) or (num_samples, num_features)
        Returns:
            cluster_ids <torch.Tensor>: (batch_size, num_samples) or (num_samples,)
        """
        n_dims = data.dim()

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

        normalized_data = data / torch.norm(data, dim=2, keepdim=True)

        if self.training:
            if self.cluster_ids is None or self.centroids is None:
                cluster_ids, centroids = self._init_kmeans(normalized_data)
                self.cluster_ids, self.centroids = cluster_ids, centroids

            if iteration is not None:
                for idx in range(iteration):
                    self.cluster_ids, self.centroids = self.update_once(normalized_data, cluster_ids=self.cluster_ids, centroids=self.centroids)
            else:
                while True:
                    cluster_ids, centroids = self.update_once(normalized_data, cluster_ids=self.cluster_ids, centroids=self.centroids)
                    distance = _euclid_distance(self.centroids, centroids, dim=-1)
                    distance = distance.mean().item()

                    self.cluster_ids, self.centroids = cluster_ids, centroids

                    if distance == 0:
                        break
        else:
            cluster_ids = self.infer(normalized_data)
            centroids = self.centroids

        if n_dims == 2:
            cluster_ids = cluster_ids.squeeze(dim=0) # (num_samples,)
            centroids = centroids.squeeze(dim=0) # (K, num_features)

            self.cluster_ids, self.centroids = cluster_ids, centroids

        return self.cluster_ids
//  [python]
def forward(self, data, iteration=None):
        """
        Args:
            data <torch.Tensor>: (batch_size, num_samples, num_features) or (num_samples, num_features)
        Returns:
            cluster_ids <torch.Tensor>: (batch_size, num_samples) or (num_samples,)
        """
        n_dims = data.dim()

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

        normalized_data = data / torch.linalg.norm(data, dim=2, keepdim=True)

        if self.training:
            if self.cluster_ids is None or self.centroids is None:
                cluster_ids, centroids = self._init_kmeans(normalized_data)
                self.cluster_ids, self.centroids = cluster_ids, centroids

            if iteration is not None:
                for idx in range(iteration):
                    self.cluster_ids, self.centroids = self.update_once(normalized_data, cluster_ids=self.cluster_ids, centroids=self.centroids)
            else:
                while True:
                    cluster_ids, centroids = self.update_once(normalized_data, cluster_ids=self.cluster_ids, centroids=self.centroids)
                    distance = _euclid_distance(self.centroids, centroids, dim=-1)
                    distance = distance.mean().item()

                    self.cluster_ids, self.centroids = cluster_ids, centroids

                    if distance == 0:
                        break
        else:
            cluster_ids = self.infer(normalized_data)
            centroids = self.centroids

        if n_dims == 2:
            cluster_ids = cluster_ids.squeeze(dim=0) # (num_samples,)
            centroids = centroids.squeeze(dim=0) # (K, num_features)

            self.cluster_ids, self.centroids = cluster_ids, centroids

        return self.cluster_ids
//  [python]
def get_images_color_similarity(self, inputs: Tensor,
                                    image_masks: Tensor) -> Tensor:
        """Compute the image color similarity in LAB color space."""
        assert inputs.dim() == 4
        assert inputs.size(0) == 1

        unfolded_images = unfold_wo_center(
            inputs,
            kernel_size=self.pairwise_size,
            dilation=self.pairwise_dilation)
        diff = inputs[:, :, None] - unfolded_images
        similarity = torch.exp(-torch.norm(diff, dim=1) * 0.5)

        unfolded_weights = unfold_wo_center(
            image_masks[None, None],
            kernel_size=self.pairwise_size,
            dilation=self.pairwise_dilation)
        unfolded_weights = torch.max(unfolded_weights, dim=1)[0]

        return similarity * unfolded_weights
//  [python]
def get_images_color_similarity(self, inputs: Tensor,
                                    image_masks: Tensor) -> Tensor:
        """Compute the image color similarity in LAB color space."""
        assert inputs.dim() == 4
        assert inputs.size(0) == 1

        unfolded_images = unfold_wo_center(
            inputs,
            kernel_size=self.pairwise_size,
            dilation=self.pairwise_dilation)
        diff = inputs[:, :, None] - unfolded_images
        similarity = torch.exp(-torch.linalg.norm(diff, dim=1) * 0.5)

        unfolded_weights = unfold_wo_center(
            image_masks[None, None],
            kernel_size=self.pairwise_size,
            dilation=self.pairwise_dilation)
        unfolded_weights = torch.max(unfolded_weights, dim=1)[0]

        return similarity * unfolded_weights
//  [python]
def loss1(self, pull, inspired, weights):
        return pull * torch.norm(inspired - weights)**2
//  [python]
def loss1(self, pull, inspired, weights):
        return pull * torch.linalg.norm(inspired - weights)**2
//  [python]
def regloss(self, lmbd, base_weight, base_bias=None):
        reg = lmbd * torch.norm(self.classifier.weight[:base_weight.size(0),:] - base_weight)
        if base_bias is not None:
            reg += lmbd * torch.norm(self.classifier.bias[:base_weight.size(0)] - base_bias)**2
        return reg
//  [python]
def regloss(self, lmbd, base_weight, base_bias=None):
        reg = lmbd * torch.linalg.norm(self.classifier.weight[:base_weight.size(0),:] - base_weight)
        if base_bias is not None:
            reg += lmbd * torch.linalg.norm(self.classifier.bias[:base_weight.size(0)] - base_bias)**2
        return reg
//  [python]
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    """
    Taken from https://github.com/mkocabas/VIBE/blob/master/lib/utils/geometry.py
    Calculates the rotation matrices for a batch of rotation vectors
    - param rot_vecs: torch.tensor (N, 3) array of N axis-angle vectors
    - returns R: torch.tensor (N, 3, 3) rotation matrices
    """
    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view(
        (batch_size, 3, 3)
    )

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat
//  [python]
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    """
    Taken from https://github.com/mkocabas/VIBE/blob/master/lib/utils/geometry.py
    Calculates the rotation matrices for a batch of rotation vectors
    - param rot_vecs: torch.tensor (N, 3) array of N axis-angle vectors
    - returns R: torch.tensor (N, 3, 3) rotation matrices
    """
    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.linalg.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view(
        (batch_size, 3, 3)
    )

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat
//  [python]
def reps_l2_sim(sent1_reps: torch.Tensor, sent2_reps: torch.Tensor) -> torch.Tensor:
    """
    calculate representation L2 similarity
    :param sent1_reps: (N, sent1_len, reps_dim)
    :param sent2_reps: (N, sent2_len, reps_dim)
    :return: (N, sent1_len, sent2_len)
    """
    sent1_len = sent1_reps.shape[-2]
    sent2_len = sent2_reps.shape[-2]

    expand_shape1 = list(sent2_reps.shape)
    expand_shape1.insert(2, sent2_len)
    expand_shape2 = list(sent2_reps.shape)
    expand_shape2.insert(1, sent1_len)

    # shape: (N, seq_len1, seq_len2, emb_dim)
    expand_reps1 = sent1_reps.unsqueeze(2).expand(expand_shape1)
    expand_reps2 = sent2_reps.unsqueeze(1).expand(expand_shape2)

    # shape: (N, seq_len1, seq_len2)
    sim = torch.norm(expand_reps1 - expand_reps2, dim=-1, p=2)
    return -sim
//  [python]
def reps_l2_sim(sent1_reps: torch.Tensor, sent2_reps: torch.Tensor) -> torch.Tensor:
    """
    calculate representation L2 similarity
    :param sent1_reps: (N, sent1_len, reps_dim)
    :param sent2_reps: (N, sent2_len, reps_dim)
    :return: (N, sent1_len, sent2_len)
    """
    sent1_len = sent1_reps.shape[-2]
    sent2_len = sent2_reps.shape[-2]

    expand_shape1 = list(sent2_reps.shape)
    expand_shape1.insert(2, sent2_len)
    expand_shape2 = list(sent2_reps.shape)
    expand_shape2.insert(1, sent1_len)

    # shape: (N, seq_len1, seq_len2, emb_dim)
    expand_reps1 = sent1_reps.unsqueeze(2).expand(expand_shape1)
    expand_reps2 = sent2_reps.unsqueeze(1).expand(expand_shape2)

    # shape: (N, seq_len1, seq_len2)
    sim = torch.linalg.norm(expand_reps1 - expand_reps2, dim=-1, ord=2)
    return -sim
//  [python]
def vectors2rotmat(vec_rest, vec_final, dtype):
    # type: (Tensor, Tensor, int) -> Tensor
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc
//  [python]
def vectors2rotmat(vec_rest, vec_final, dtype):
    # type: (Tensor, Tensor, int) -> Tensor
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.linalg.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.linalg.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc
//  [python]
def compute_error_accel_lsp14(joints_gt, joints_pred):
    """
    Computes acceleration error:
        1/(n-2) \sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}
    Note that for each frame that is not visible, three entries in the
    acceleration error should be zero'd out.
    Args:
        joints_gt (Nx14x3).
        joints_pred (Nx14x3).
        vis (N).
    Returns:
        error_accel (N-2).
    """
    # (N-2)x14x3
    accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
    accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]

    vis_mask = (joints_gt != -2.).sum(-1) >2
    vis = vis_mask[:-2] * vis_mask[1:-1] * vis_mask[2:]

    normed = torch.norm(accel_pred - accel_gt, p=2, dim=2)

    return normed[vis]
//  [python]
def compute_error_accel_lsp14(joints_gt, joints_pred):
    """
    Computes acceleration error:
        1/(n-2) \sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}
    Note that for each frame that is not visible, three entries in the
    acceleration error should be zero'd out.
    Args:
        joints_gt (Nx14x3).
        joints_pred (Nx14x3).
        vis (N).
    Returns:
        error_accel (N-2).
    """
    # (N-2)x14x3
    accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
    accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]

    vis_mask = (joints_gt != -2.).sum(-1) >2
    vis = vis_mask[:-2] * vis_mask[1:-1] * vis_mask[2:]

    normed = torch.linalg.norm(accel_pred - accel_gt, p=2, dim=2)

    return normed[vis]
//  [python]
def reglossnovel(self, lmbd, novel_weight, novel_bias=None):
        rng1, rng2 = self.num_classes, self.num_classes + novel_weight.size(0)
        reg = lmbd * torch.norm(self.classifier.weight[rng1:rng2, :] - novel_weight) #**2??
        if novel_bias is not None:
            reg += lmbd * torch.norm(self.classifier.bias[rng1:rng2, :] - novel_bias)**2
        return reg
//  [python]
def reglossnovel(self, lmbd, novel_weight, novel_bias=None):
        rng1, rng2 = self.num_classes, self.num_classes + novel_weight.size(0)
        reg = lmbd * torch.linalg.norm(self.classifier.weight[rng1:rng2, :] - novel_weight) #**2??
        if novel_bias is not None:
            reg += lmbd * torch.linalg.norm(self.classifier.bias[rng1:rng2, :] - novel_bias)**2
        return reg
//  [python]
def get_sing_loss(w,loss_value,error=None,args=None):
    sp=w.shape
    def normed(mat):
        normer = torch.norm(mat)
        if normer !=0:
            return mat / normer
        print('Normer = 0 !!!')
        return 0

    normed_w = normed(w.view(sp[0],-1))
    if isinstance(normed_w,int) and normed_w == 0:
        return 0.,0
    try:
        U,S,V = torch.svd(normed_w.detach())
        # flops calc for reb
        m, n = normed_w.shape
        args.svd_flops['solve'] += (2*m*n*n+11*n*n*n)
        args.svd_flops['iter'] += (4*m*n*n-4/3*n*n*n)
        args.svd_flops['min_solve'] += (2*m*n*n+11*n*n*n)
        args.svd_flops['min_iter'] += (4*m*n*n-4/3*n*n*n)
    except:
        rand_num = int(time.time())
        file_name = f'results/err{rand_num}.pkl'
        print(f'Error Detected ... Error matrix saved in {file_name}!')
        torch.save(normed_w.detach(),file_name)
        return str(rand_num),0
    
    k = bsearch(normed_w,U,S,V,0,normed_w.size(0), loss_value=loss_value)
    assert k <= sp[0]
    w_approx = U[:,:k] @ torch.diag(S)[:k, :k]@V.t()[:k,:]
    loss =  -F.mse_loss(normed_w,w_approx,reduction='sum')
    if error is None:
        return loss,k
    else:
        l = (loss - error)
        if l < 0:
            return  l ** 2,k
        else:
            return 0.,k
//  [python]
def get_sing_loss(w,loss_value,error=None,args=None):
    sp=w.shape
    def normed(mat):
        normer = torch.linalg.norm(mat)
        if normer !=0:
            return mat / normer
        print('Normer = 0 !!!')
        return 0

    normed_w = normed(w.view(sp[0],-1))
    if isinstance(normed_w,int) and normed_w == 0:
        return 0.,0
    try:
        U,S,V = torch.svd(normed_w.detach())
        # flops calc for reb
        m, n = normed_w.shape
        args.svd_flops['solve'] += (2*m*n*n+11*n*n*n)
        args.svd_flops['iter'] += (4*m*n*n-4/3*n*n*n)
        args.svd_flops['min_solve'] += (2*m*n*n+11*n*n*n)
        args.svd_flops['min_iter'] += (4*m*n*n-4/3*n*n*n)
    except:
        rand_num = int(time.time())
        file_name = f'results/err{rand_num}.pkl'
        print(f'Error Detected ... Error matrix saved in {file_name}!')
        torch.save(normed_w.detach(),file_name)
        return str(rand_num),0
    
    k = bsearch(normed_w,U,S,V,0,normed_w.size(0), loss_value=loss_value)
    assert k <= sp[0]
    w_approx = U[:,:k] @ torch.diag(S)[:k, :k]@V.t()[:k,:]
    loss =  -F.mse_loss(normed_w,w_approx,reduction='sum')
    if error is None:
        return loss,k
    else:
        l = (loss - error)
        if l < 0:
            return  l ** 2,k
        else:
            return 0.,k
//  [python]
def isin(self,xnext):
		for v in self.x:
			if torch.norm(v-xnext) < self.epsilon:
				return True

		return False
//  [python]
def isin(self,xnext):
		for v in self.x:
			if torch.linalg.norm(v-xnext) < self.epsilon:
				return True

		return False
//  [python]
def trs_decomp(A):
    if has_torch(A):
        s_vec = torch.norm(A[:3, :3], dim=0)
    else:
        s_vec = np.linalg.norm(A[:3, :3], axis=0)
    R = A[:3, :3] / s_vec
    t = A[:3, 3]
    return t, R, s_vec
//  [python]
def trs_decomp(A):
    if has_torch(A):
        s_vec = torch.linalg.norm(A[:3, :3], dim=0)
    else:
        s_vec = np.linalg.norm(A[:3, :3], axis=0)
    R = A[:3, :3] / s_vec
    t = A[:3, 3]
    return t, R, s_vec
//  [python]
def isin_hal(self,xnext,x_hal):
		for v in x_hal:
			if torch.norm(v-xnext) < self.epsilon:
				return True
//  [python]
def isin_hal(self,xnext,x_hal):
		for v in x_hal:
			if torch.linalg.norm(v-xnext) < self.epsilon:
				return True
//  [python]
def test_norm(self, device):
        # full reduction
        x = torch.randn(25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = x.norm(p).item()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res, expected, "full reduction failed for {}-norm".format(p))

        # one dimension
        x = torch.randn(25, 25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = x.norm(p, 1).cpu().numpy()
            expected = np.linalg.norm(xn, p, 1)
            self.assertEqual(res.shape, expected.shape)
            self.assertTrue(np.allclose(res, expected), "dim reduction failed for {}-norm".format(p))

        # matrix norm
        for p in ['fro', 'nuc']:
            res = x.norm(p).cpu().numpy()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res.shape, expected.shape)
            self.assertTrue(np.allclose(res, expected), "dim reduction failed for {}-norm".format(p))

        # larger tensor sanity check
        self.assertEqual(2 * torch.norm(torch.ones(10000)), torch.norm(torch.ones(40000)))
//  [python]
def test_norm(self, device):
        # full reduction
        x = torch.randn(25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = torch.linalg.norm(x, p).item()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res, expected, "full reduction failed for {}-norm".format(p))

        # one dimension
        x = torch.randn(25, 25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = torch.linalg.norm(x, p, 1).cpu().numpy()
            expected = np.linalg.norm(xn, p, 1)
            self.assertEqual(res.shape, expected.shape)
            self.assertTrue(np.allclose(res, expected), "dim reduction failed for {}-norm".format(p))

        # matrix norm
        for p in ['fro', 'nuc']:
            res = torch.linalg.norm(x, p).cpu().numpy()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res.shape, expected.shape)
            self.assertTrue(np.allclose(res, expected), "dim reduction failed for {}-norm".format(p))

        # larger tensor sanity check
        self.assertEqual(2 * torch.linalg.norm(torch.ones(10000)), torch.linalg.norm(torch.ones(40000)))
//  [python]
def forward(self, g, h):
        feat = self.feat_gc(
            g, h
        )  # size = (sum_N, F_out), sum_N is num of nodes in this batch
        device = feat.device
        assign_tensor = self.pool_gc(
            g, h
        )  # size = (sum_N, N_a), N_a is num of nodes in pooled graph.
        assign_tensor = F.softmax(assign_tensor, dim=1)
        assign_tensor = torch.split(assign_tensor, g.batch_num_nodes().tolist())
        assign_tensor = torch.block_diag(
            *assign_tensor
        )  # size = (sum_N, batch_size * N_a)

        h = torch.matmul(torch.t(assign_tensor), feat)
        adj = g.adj_external(transpose=True, ctx=device)
        adj_new = torch.sparse.mm(adj, assign_tensor)
        adj_new = torch.mm(torch.t(assign_tensor), adj_new)

        if self.link_pred:
            current_lp_loss = torch.norm(
                adj.to_dense() - torch.mm(assign_tensor, torch.t(assign_tensor))
            ) / np.power(g.num_nodes(), 2)
            self.loss_log["LinkPredLoss"] = current_lp_loss

        for loss_layer in self.reg_loss:
            loss_name = str(type(loss_layer).__name__)
            self.loss_log[loss_name] = loss_layer(adj, adj_new, assign_tensor)

        return adj_new, h
//  [python]
def forward(self, g, h):
        feat = self.feat_gc(
            g, h
        )  # size = (sum_N, F_out), sum_N is num of nodes in this batch
        device = feat.device
        assign_tensor = self.pool_gc(
            g, h
        )  # size = (sum_N, N_a), N_a is num of nodes in pooled graph.
        assign_tensor = F.softmax(assign_tensor, dim=1)
        assign_tensor = torch.split(assign_tensor, g.batch_num_nodes().tolist())
        assign_tensor = torch.block_diag(
            *assign_tensor
        )  # size = (sum_N, batch_size * N_a)

        h = torch.matmul(torch.t(assign_tensor), feat)
        adj = g.adj_external(transpose=True, ctx=device)
        adj_new = torch.sparse.mm(adj, assign_tensor)
        adj_new = torch.mm(torch.t(assign_tensor), adj_new)

        if self.link_pred:
            current_lp_loss = torch.linalg.norm(
                adj.to_dense() - torch.mm(assign_tensor, torch.t(assign_tensor))
            ) / np.power(g.num_nodes(), 2)
            self.loss_log["LinkPredLoss"] = current_lp_loss

        for loss_layer in self.reg_loss:
            loss_name = str(type(loss_layer).__name__)
            self.loss_log[loss_name] = loss_layer(adj, adj_new, assign_tensor)

        return adj_new, h
//  [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])  # Bx1xL
    y = y.view(1, -1, x.shape[-1])  # 1xBxL
    dist = th.norm(x - y, dim=2) / (num_grid_per_dim * raw_length_scale)
    return signal_variance * th.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])  # Bx1xL
    y = y.view(1, -1, x.shape[-1])  # 1xBxL
    dist = th.linalg.norm(x - y, dim=2) / (num_grid_per_dim * raw_length_scale)
    return signal_variance * th.exp(-dist)
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    # type: (Tensor, Tensor, Tensor, Tensor, int) -> Tensor
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    assert children[0] == 3
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, int(children[0])].clone()
    spine_rest_loc = rel_rest_pose[:, int(children[0])].clone()
    spine_norm = torch.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine)
    center_final_loc = torch.tensor(0, dtype=dtype, device=rel_pose_skeleton.device)
    center_rest_loc = torch.tensor(0, dtype=dtype, device=rel_pose_skeleton.device)
    for child in pelvis_child:
        if child == int(children[0]):
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()
    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.norm(center_rest_loc, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    # type: (Tensor, Tensor, Tensor, Tensor, int) -> Tensor
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    assert children[0] == 3
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, int(children[0])].clone()
    spine_rest_loc = rel_rest_pose[:, int(children[0])].clone()
    spine_norm = torch.linalg.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine)
    center_final_loc = torch.tensor(0, dtype=dtype, device=rel_pose_skeleton.device)
    center_rest_loc = torch.tensor(0, dtype=dtype, device=rel_pose_skeleton.device)
    for child in pelvis_child:
        if child == int(children[0]):
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()
    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.linalg.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.linalg.norm(center_rest_loc, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))
    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def update(self):
        """Set and update iteration variables.
        """
        if self.ivars['istep'] == 0:
            X = torch.randn_like(self.X)
            X_norm = float(torch.norm(X))
            iX_norm = X_norm ** -1
            A_norm = float(torch.norm(_utils.matmul(self.A, X))) * iX_norm
            B_norm = float(torch.norm(_utils.matmul(self.B, X))) * iX_norm
            self.fvars['X_norm'] = X_norm
            self.fvars['A_norm'] = A_norm
            self.fvars['B_norm'] = B_norm
            self.ivars['iterations_left'] = self.iparams['niter']
            self.ivars['converged_count'] = 0
            self.ivars['converged_end'] = 0

        if self.method == 'ortho':
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars['iterations_left'] = self.ivars['iterations_left'] - 1
        self.ivars['istep'] = self.ivars['istep'] + 1
//  [python]
def update(self):
        """Set and update iteration variables.
        """
        if self.ivars['istep'] == 0:
            X = torch.randn_like(self.X)
            X_norm = float(torch.linalg.norm(X))
            iX_norm = X_norm ** -1
            A_norm = float(torch.linalg.norm(_utils.matmul(self.A, X))) * iX_norm
            B_norm = float(torch.linalg.norm(_utils.matmul(self.B, X))) * iX_norm
            self.fvars['X_norm'] = X_norm
            self.fvars['A_norm'] = A_norm
            self.fvars['B_norm'] = B_norm
            self.ivars['iterations_left'] = self.iparams['niter']
            self.ivars['converged_count'] = 0
            self.ivars['converged_end'] = 0

        if self.method == 'ortho':
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars['iterations_left'] = self.ivars['iterations_left'] - 1
        self.ivars['istep'] = self.ivars['istep'] + 1
//  [python]
def step_correct(
        self,
        model_output: torch.FloatTensor,
        sample: torch.FloatTensor,
        generator: Optional[torch.Generator] = None,
        return_dict: bool = True,
    ) -> Union[SchedulerOutput, Tuple]:
        """
        Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after
        making the prediction for the previous timestep.

        Args:
            model_output (`torch.FloatTensor`):
                The direct output from learned diffusion model.
            sample (`torch.FloatTensor`):
                A current instance of a sample created by the diffusion process.
            generator (`torch.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

        Returns:
            [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
                is returned where the first element is the sample tensor.

        """
        if self.timesteps is None:
            raise ValueError(
                "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
            )

        # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
        # sample noise for correction
        noise = randn_tensor(sample.shape, layout=sample.layout, generator=generator).to(sample.device)

        # compute step size from the model_output, the noise, and the snr
        grad_norm = torch.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean()
        noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
        step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
        step_size = step_size * torch.ones(sample.shape[0]).to(sample.device)
        # self.repeat_scalar(step_size, sample.shape[0])

        # compute corrected sample: model_output term and noise term
        step_size = step_size.flatten()
        while len(step_size.shape) < len(sample.shape):
            step_size = step_size.unsqueeze(-1)
        prev_sample_mean = sample + step_size * model_output
        prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise

        if not return_dict:
            return (prev_sample,)

        return SchedulerOutput(prev_sample=prev_sample)