jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def rel(self, x, y, mask=None):
        """
        Args:
            x, y: both have shape of [B, -1]
        """
        num_examples = x.size()[0]

        if mask is None:
            diff_norms = torch.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), self.p, 1)
            y_norms = torch.norm(y.reshape(num_examples,-1), self.p, 1)
        else:
            mask = mask.view(num_examples, -1).float()
            diff_norms = torch.norm((x.reshape(num_examples,-1) - y.reshape(num_examples,-1)) * mask, self.p, 1)
            y_norms = torch.norm(y.reshape(num_examples,-1) * mask + (1 - mask) * 1e-6, self.p, 1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms/y_norms)
            else:
                return torch.sum(diff_norms/y_norms)

        return diff_norms/y_norms
//  [python]
def rel(self, x, y, mask=None):
        """
        Args:
            x, y: both have shape of [B, -1]
        """
        num_examples = x.size()[0]

        if mask is None:
            diff_norms = torch.linalg.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), ord=self.p, dim=1)
            y_norms = torch.linalg.norm(y.reshape(num_examples,-1), ord=self.p, dim=1)
        else:
            mask = mask.view(num_examples, -1).float()
            diff_norms = torch.linalg.norm((x.reshape(num_examples,-1) - y.reshape(num_examples,-1)) * mask, ord=self.p, dim=1)
            y_norms = torch.linalg.norm(y.reshape(num_examples,-1) * mask + (1 - mask) * 1e-6, ord=self.p, dim=1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms/y_norms)
            else:
                return torch.sum(diff_norms/y_norms)

        return diff_norms/y_norms
//  [python]
def transformation_residuals(x1, x2, R, t):
    """
    Computer the pointwise residuals based on the estimated transformation paramaters
    
    Args:
        x1  (torch array): points of the first point cloud [b,n,3]
        x2  (torch array): points of the second point cloud [b,n,3]
        R   (torch array): estimated rotation matrice [b,3,3]
        t   (torch array): estimated translation vectors [b,3,1]
    Returns:
        res (torch array): pointwise residuals (Eucledean distance) [b,n,1]
    """
    x2_reconstruct = torch.matmul(R, x1.transpose(1, 2)) + t 

    res = torch.norm(x2_reconstruct.transpose(1, 2) - x2, dim=2)

    return res
//  [python]
def transformation_residuals(x1, x2, R, t):
    """
    Computer the pointwise residuals based on the estimated transformation paramaters
    
    Args:
        x1  (torch array): points of the first point cloud [b,n,3]
        x2  (torch array): points of the second point cloud [b,n,3]
        R   (torch array): estimated rotation matrice [b,3,3]
        t   (torch array): estimated translation vectors [b,3,1]
    Returns:
        res (torch array): pointwise residuals (Eucledean distance) [b,n,1]
    """
    x2_reconstruct = torch.matmul(R, x1.transpose(1, 2)) + t 

    res = torch.linalg.norm(x2_reconstruct.transpose(1, 2) - x2, dim=2, keepdim=True)

    return res
//  [python]
def line_search(update, x0, g0, g, nstep=0, on=True):
    """
    `update` is the propsoed direction of update.

    Code adapted from scipy.
    """
    tmp_s = [0]
    tmp_g0 = [g0]
    tmp_phi = [torch.norm(g0)**2]
    s_norm = torch.norm(x0) / torch.norm(update)

    def phi(s, store=True):
        if s == tmp_s[0]:
            return tmp_phi[0]    # If the step size is so small... just return something
        x_est = x0 + s * update
        g0_new = g(x_est)
        phi_new = _safe_norm(g0_new)**2
        if store:
            tmp_s[0] = s
            tmp_g0[0] = g0_new
            tmp_phi[0] = phi_new
        return phi_new

    if on:
        s, phi1, ite = scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0], amin=1e-2)
    if (not on) or s is None:
        s = 1.0
        ite = 0

    x_est = x0 + s * update
    if s == tmp_s[0]:
        g0_new = tmp_g0[0]
    else:
        g0_new = g(x_est)
    return x_est, g0_new, x_est - x0, g0_new - g0, ite
//  [python]
def line_search(update, x0, g0, g, nstep=0, on=True):
    """
    `update` is the propsoed direction of update.

    Code adapted from scipy.
    """
    tmp_s = [0]
    tmp_g0 = [g0]
    tmp_phi = [torch.linalg.norm(g0)**2]
    s_norm = torch.linalg.norm(x0) / torch.linalg.norm(update)

    def phi(s, store=True):
        if s == tmp_s[0]:
            return tmp_phi[0]    # If the step size is so small... just return something
        x_est = x0 + s * update
        g0_new = g(x_est)
        phi_new = _safe_norm(g0_new)**2
        if store:
            tmp_s[0] = s
            tmp_g0[0] = g0_new
            tmp_phi[0] = phi_new
        return phi_new

    if on:
        s, phi1, ite = scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0], amin=1e-2)
    if (not on) or s is None:
        s = 1.0
        ite = 0

    x_est = x0 + s * update
    if s == tmp_s[0]:
        g0_new = tmp_g0[0]
    else:
        g0_new = g(x_est)
    return x_est, g0_new, x_est - x0, g0_new - g0, ite
//  [python]
def batch_rodrigues(param):
    #param N x 3
    batch_size = param.shape[0]
    #沿第二维（3个数）进行求二次范数：||x||，下面就是进行标准化，每三个数除以他们的范数。
    l1norm = torch.norm(param + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(param, angle)
    angle = angle * 0.5
    #上面算出的是一个向量的长度：sqrt(x**2+y**2+z**2)/2,所以这个长度的的cos
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    #用四元组表示三维旋转，有时间看一下×××××××××
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)

    return quat2mat(quat)
//  [python]
def batch_rodrigues(param):
    #param N x 3
    batch_size = param.shape[0]
    #沿第二维（3个数）进行求二次范数：||x||，下面就是进行标准化，每三个数除以他们的范数。
    l1norm = torch.linalg.norm(param + 1e-8, ord = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(param, angle)
    angle = angle * 0.5
    #上面算出的是一个向量的长度：sqrt(x**2+y**2+z**2)/2,所以这个长度的的cos
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    #用四元组表示三维旋转，有时间看一下×××××××××
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)

    return quat2mat(quat)
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})

    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]

    Returns:
        te (torch tensor): translation error [b,1]

    """
    return torch.norm(t1-t2, dim=(-2,-1))
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})

    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]

    Returns:
        te (torch tensor): translation error [b,1]

    """
    return torch.linalg.norm(t1-t2, dim=(-2,-1))
//  [python]
def get_regularization(model_list, reg_type_core):
    """Get regularization.

    Args:
        reg_type_core, Choose from:
            "None": no regularization
            "l1": L1 regularization
            "l2": L2 regularization
            "nuc": nuclear regularization
            "fro": Frobenius norm
            "snr": spectral regularization
            "snn": spectral normalization
            "SINDy": L1 regularization on the coefficient of SINDy
            "Hall": regularize all the elements of Hessian
            "Hoff": regularize the off-diagonal elements of Hessian
            "Hdiag": regularize the diagonal elements of Hessian

    Returns:
        reg: computed regularization.
    """
    if reg_type_core in ["None", "snn"]:
        return 0
    else:
        List = []
        if reg_type_core in ["l1", "l2", "nuc", "fro"]:
            for model in model_list:
                for param_key, param in model.named_parameters():
                    if "weight" in param_key and param.requires_grad:
                        if reg_type_core in ["nuc", "fro"]:
                            norm = torch.norm(param, reg_type_core)
                        elif reg_type_core == "l1":
                            norm = param.abs().sum()
                        elif reg_type_core == "l2":
                            norm = param.square().sum()
                        else:
                            raise
                        List.append(norm)
        elif reg_type_core == "snr":
            for model in model_list:
                for module in model.modules():
                    if module.__class__.__name__ == "SpectralNormReg" and hasattr(module, "snreg"):
                        List.append(module.snreg)
        elif reg_type_core == "sindy":
            for model in model_list:
                for module in model.modules():
                    if module.__class__.__name__ == "SINDy":
                        List.append(module.weight.abs().sum())
        elif reg_type_core in ["Hall", "Hoff", "Hdiag"]:
            for model in model_list:
                if hasattr(model, "Hreg"):
                    List.append(model.Hreg)
        else:
            raise Exception("reg_type_core {} is not valid!".format(reg_type_core))
        if len(List) > 0:
            reg = torch.stack(List).sum()
        else:
            reg = 0
        return reg
//  [python]
def get_regularization(model_list, reg_type_core):
    """Get regularization.

    Args:
        reg_type_core, Choose from:
            "None": no regularization
            "l1": L1 regularization
            "l2": L2 regularization
            "nuc": nuclear regularization
            "fro": Frobenius norm
            "snr": spectral regularization
            "snn": spectral normalization
            "SINDy": L1 regularization on the coefficient of SINDy
            "Hall": regularize all the elements of Hessian
            "Hoff": regularize the off-diagonal elements of Hessian
            "Hdiag": regularize the diagonal elements of Hessian

    Returns:
        reg: computed regularization.
    """
    if reg_type_core in ["None", "snn"]:
        return 0
    else:
        List = []
        if reg_type_core in ["l1", "l2", "nuc", "fro"]:
            for model in model_list:
                for param_key, param in model.named_parameters():
                    if "weight" in param_key and param.requires_grad:
                        if reg_type_core in ["nuc", "fro"]:
                            norm = torch.linalg.norm(param, ord=reg_type_core)
                        elif reg_type_core == "l1":
                            norm = param.abs().sum()
                        elif reg_type_core == "l2":
                            norm = param.square().sum()
                        else:
                            raise
                        List.append(norm)
        elif reg_type_core == "snr":
            for model in model_list:
                for module in model.modules():
                    if module.__class__.__name__ == "SpectralNormReg" and hasattr(module, "snreg"):
                        List.append(module.snreg)
        elif reg_type_core == "sindy":
            for model in model_list:
                for module in model.modules():
                    if module.__class__.__name__ == "SINDy":
                        List.append(module.weight.abs().sum())
        elif reg_type_core in ["Hall", "Hoff", "Hdiag"]:
            for model in model_list:
                if hasattr(model, "Hreg"):
                    List.append(model.Hreg)
        else:
            raise Exception("reg_type_core {} is not valid!".format(reg_type_core))
        if len(List) > 0:
            reg = torch.stack(List).sum()
        else:
            reg = 0
        return reg
//  [python]
def distances(self, positions, index: Optional[torch.Tensor] = None):
        """Calculate the distances with tuples, starting with the offset used."""
        if index is None:
            index = self.build_indices(positions.size(1), 2, self.distances_offset)
        doubles = positions[:, index, :]

        return torch.norm(doubles[..., 1, :] - doubles[..., 0, :], dim=-1)
//  [python]
def distances(self, positions, index: Optional[torch.Tensor] = None):
        """Calculate the distances with tuples, starting with the offset used."""
        if index is None:
            index = self.build_indices(positions.size(1), 2, self.distances_offset)
        doubles = positions[:, index, :]

        return torch.linalg.norm(doubles[..., 1, :] - doubles[..., 0, :], dim=-1)
//  [python]
def extract_bd_theta(positions: torch.FloatTensor):
    """Extracts the bond distance and angle in radians for a given tetrad structure in torch
    tensor.

    Args:
        positions (torch.FloatTensor): A tensor representing the positions of atoms in a tetrad structure.
            The structure is given by (a,b,c,d), where d is a branching atom.
            Dimensions: K x B x 4 x 3.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: The calculated bond distance and angle in radians.
    """
    a = positions[:, :, 0, :]
    b = positions[:, :, 1, :]
    c = positions[:, :, 2, :]
    d = positions[:, :, 3, :]

    # calculate the main vectors
    ba = a - b
    bc = c - b
    bd = d - b

    cross_product = torch.cross(bc, ba, dim=-1)
    norm = torch.norm(cross_product, dim=-1, keepdim=False)
    bd_norm = torch.norm(bd, dim=-1, keepdim=False)

    # calculate the angle between the normal vector and bd vector
    dot_prod = (cross_product * bd).sum(dim=-1)
    cos_angle = dot_prod / (norm * bd_norm)
    angle_radians = torch.arccos(cos_angle)

    return bd_norm, angle_radians
//  [python]
def extract_bd_theta(positions: torch.FloatTensor):
    """Extracts the bond distance and angle in radians for a given tetrad structure in torch
    tensor.

    Args:
        positions (torch.FloatTensor): A tensor representing the positions of atoms in a tetrad structure.
            The structure is given by (a,b,c,d), where d is a branching atom.
            Dimensions: K x B x 4 x 3.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: The calculated bond distance and angle in radians.
    """
    a = positions[:, :, 0, :]
    b = positions[:, :, 1, :]
    c = positions[:, :, 2, :]
    d = positions[:, :, 3, :]

    # calculate the main vectors
    ba = a - b
    bc = c - b
    bd = d - b

    cross_product = torch.cross(bc, ba, dim=-1)
    norm = torch.linalg.norm(cross_product, dim=-1, keepdim=False)
    bd_norm = torch.linalg.norm(bd, dim=-1, keepdim=False)

    # calculate the angle between the normal vector and bd vector
    dot_prod = (cross_product * bd).sum(dim=-1)
    cos_angle = dot_prod / (norm * bd_norm)
    angle_radians = torch.arccos(cos_angle)

    return bd_norm, angle_radians
//  [python]
def quat_error_magnitude(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
    """Computes the rotation difference between two quaternions.

    Args:
        q1: The first quaternion in (w, x, y, z). Shape is (..., 4).
        q2: The second quaternion in (w, x, y, z). Shape is (..., 4).

    Returns:
        Angular error between input quaternions in radians.
    """
    quat_diff = quat_mul(q1, quat_conjugate(q2))
    return torch.norm(axis_angle_from_quat(quat_diff), dim=-1)
//  [python]
def quat_error_magnitude(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
    """Computes the rotation difference between two quaternions.

    Args:
        q1: The first quaternion in (w, x, y, z). Shape is (..., 4).
        q2: The second quaternion in (w, x, y, z). Shape is (..., 4).

    Returns:
        Angular error between input quaternions in radians.
    """
    quat_diff = quat_mul(q1, quat_conjugate(q2))
    return torch.linalg.norm(axis_angle_from_quat(quat_diff), dim=-1)
//  [python]
def _get_binding_edges(self, X, S, inter_edges):
        atom_pos = self.aa_feature._construct_atom_pos(S)
        src_dst = inter_edges.T
        dist = X[src_dst]  # [Ef, 2, n_channel, 3]
        dist = dist[:, 0].unsqueeze(2) - dist[:, 1].unsqueeze(1)  # [Ef, n_channel, n_channel, 3]
        dist = torch.norm(dist, dim=-1)  # [Ef, n_channel, n_channel]
        pos_pad = atom_pos[src_dst] == self.aa_feature.atom_pos_pad_idx # [Ef, 2, n_channel]
        pos_pad = torch.logical_or(pos_pad[:, 0].unsqueeze(2), pos_pad[:, 1].unsqueeze(1))  # [Ef, n_channel, n_channel]
        dist = dist + pos_pad * 1e10  # [Ef, n_channel, n_channel]
        dist = torch.min(dist.reshape(dist.shape[0], -1), dim=1)[0]  # [Ef]
        is_binding = dist <= self.bind_dist_cutoff
        return src_dst[is_binding].T
//  [python]
def _get_binding_edges(self, X, S, inter_edges):
        atom_pos = self.aa_feature._construct_atom_pos(S)
        src_dst = inter_edges.T
        dist = X[src_dst]  # [Ef, 2, n_channel, 3]
        dist = dist[:, 0].unsqueeze(2) - dist[:, 1].unsqueeze(1)  # [Ef, n_channel, n_channel, 3]
        dist = torch.linalg.norm(dist, dim=-1)  # [Ef, n_channel, n_channel]
        pos_pad = atom_pos[src_dst] == self.aa_feature.atom_pos_pad_idx # [Ef, 2, n_channel]
        pos_pad = torch.logical_or(pos_pad[:, 0].unsqueeze(2), pos_pad[:, 1].unsqueeze(1))  # [Ef, n_channel, n_channel]
        dist = dist + pos_pad * 1e10  # [Ef, n_channel, n_channel]
        dist = torch.min(dist.reshape(dist.shape[0], -1), dim=1)[0]  # [Ef]
        is_binding = dist <= self.bind_dist_cutoff
        return src_dst[is_binding].T
//  [python]
def quat_box_minus(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
    """The box-minus operator (quaternion difference) between two quaternions.

    Args:
        q1: The first quaternion in (w, x, y, z). Shape is (N, 4).
        q2: The second quaternion in (w, x, y, z). Shape is (N, 4).

    Returns:
        The difference between the two quaternions. Shape is (N, 3).

    Reference:
        https://docs.leggedrobotics.com/kindr/cheatsheet_latest.pdf
    """
    quat_diff = quat_mul(q1, quat_conjugate(q2))  # q1 * q2^-1
    re = quat_diff[:, 0]  # real part, q = [w, x, y, z] = [re, im]
    im = quat_diff[:, 1:]  # imaginary part
    norm_im = torch.norm(im, dim=1)
    scale = 2.0 * torch.where(norm_im > 1.0e-7, torch.atan2(norm_im, re) / norm_im, torch.sign(re))
    return scale.unsqueeze(-1) * im
//  [python]
def quat_box_minus(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
    """The box-minus operator (quaternion difference) between two quaternions.

    Args:
        q1: The first quaternion in (w, x, y, z). Shape is (N, 4).
        q2: The second quaternion in (w, x, y, z). Shape is (N, 4).

    Returns:
        The difference between the two quaternions. Shape is (N, 3).

    Reference:
        https://docs.leggedrobotics.com/kindr/cheatsheet_latest.pdf
    """
    quat_diff = quat_mul(q1, quat_conjugate(q2))  # q1 * q2^-1
    re = quat_diff[:, 0]  # real part, q = [w, x, y, z] = [re, im]
    im = quat_diff[:, 1:]  # imaginary part
    norm_im = torch.linalg.norm(im, dim=1)
    scale = 2.0 * torch.where(norm_im > 1.0e-7, torch.atan2(norm_im, re) / norm_im, torch.sign(re))
    return scale.unsqueeze(-1) * im
//  [python]
def calculate_bd(self, positions: torch.Tensor, bd_norm: torch.Tensor, theta: torch.Tensor):
        """Calculates the branched atom distance for a given structure.

        Args:
            positions (torch.Tensor): A tensor representing the positions of atoms.
                Dimensions: K x B x 3 x 3, with K as the batch size, B the number of structures in the batch, and 3x3 as the position vectors.
            bd_norm (torch.Tensor): A tensor representing the norms corresponding to the branched atom distance. Vector based, size N x 1.
            theta (torch.Tensor): A tensor representing the angle in degrees. Vector based, size N x 1.

        Returns:
            torch.Tensor: The calculated branched atom distance tensor.
        """
        a = positions[:, :, 0, :]
        b = positions[:, :, 1, :]
        c = positions[:, :, 2, :]

        # calculate the main vectors
        ba = a - b
        bc = c - b

        centroid = torch.mean(torch.stack([a, c]), dim=0)
        centroid_b = b - centroid  #

        a_prime = a + centroid_b
        c_prime = c + centroid_b

        cross_product = torch.cross(bc, ba, dim=-1)
        norm = torch.norm(cross_product, dim=-1, keepdim=True)

        normalized_cross_product = (cross_product / norm) * bd_norm.unsqueeze(-1) + b

        new_bd = self.rotate_ac(a_prime, normalized_cross_product, c_prime, theta, b_only=True)

        return new_bd
//  [python]
def calculate_bd(self, positions: torch.Tensor, bd_norm: torch.Tensor, theta: torch.Tensor):
        """Calculates the branched atom distance for a given structure.

        Args:
            positions (torch.Tensor): A tensor representing the positions of atoms.
                Dimensions: K x B x 3 x 3, with K as the batch size, B the number of structures in the batch, and 3x3 as the position vectors.
            bd_norm (torch.Tensor): A tensor representing the norms corresponding to the branched atom distance. Vector based, size N x 1.
            theta (torch.Tensor): A tensor representing the angle in degrees. Vector based, size N x 1.

        Returns:
            torch.Tensor: The calculated branched atom distance tensor.
        """
        a = positions[:, :, 0, :]
        b = positions[:, :, 1, :]
        c = positions[:, :, 2, :]

        # calculate the main vectors
        ba = a - b
        bc = c - b

        centroid = torch.mean(torch.stack([a, c]), dim=0)
        centroid_b = b - centroid  #

        a_prime = a + centroid_b
        c_prime = c + centroid_b

        cross_product = torch.cross(bc, ba, dim=-1)
        norm = torch.linalg.norm(cross_product, dim=-1, keepdim=True)

        normalized_cross_product = (cross_product / norm) * bd_norm.unsqueeze(-1) + b

        new_bd = self.rotate_ac(a_prime, normalized_cross_product, c_prime, theta, b_only=True)

        return new_bd
//  [python]
def cont_to_rotmat(rot: torch.Tensor) -> torch.Tensor:
    """Convert 6D continuous rotation parameterization to roation matrix.

    Args:
        rot: continuous 6D rotation parameterization, shape (N, 6)

    Returns:
        rotation matrices, shape (N, 3, 3)
    """
    a1 = rot[:, :3]
    a2 = rot[:, 3:]
    b1 = a1 / torch.norm(a1, dim=1, keepdim=True)
    b2 = a2 - torch.sum(b1 * a2, dim=1, keepdim=True) * b1
    b2 = b2 / torch.norm(b2, dim=1, keepdim=True)
    b3 = torch.cross(b1, b2)
    return torch.cat((b1.unsqueeze(2), b2.unsqueeze(2), b3.unsqueeze(2)), dim=2)
//  [python]
def cont_to_rotmat(rot: torch.Tensor) -> torch.Tensor:
    """Convert 6D continuous rotation parameterization to roation matrix.

    Args:
        rot: continuous 6D rotation parameterization, shape (N, 6)

    Returns:
        rotation matrices, shape (N, 3, 3)
    """
    a1 = rot[:, :3]
    a2 = rot[:, 3:]
    b1 = a1 / torch.linalg.norm(a1, dim=1, keepdim=True)
    b2 = a2 - torch.sum(b1 * a2, dim=1, keepdim=True) * b1
    b2 = b2 / torch.linalg.norm(b2, dim=1, keepdim=True)
    b3 = torch.cross(b1, b2)
    return torch.cat((b1.unsqueeze(2), b2.unsqueeze(2), b3.unsqueeze(2)), dim=2)
//  [python]
def rotate_ac(self, a, b, c, theta, b_only: bool = True):
        """Performs batch rotation for the tensor, withholding K x B x 3 tensors.

        Args:
            a, b, c (torch.Tensor): Input tensors of dimensions K x B x 3, representing initial position vectors for each batch.
            theta (torch.Tensor): The angle of rotation in radians. Vector based, size N x 1.
            b_only (bool): If True, only returns tensor b_rot, otherwise also returns a and c tensors. Default is True.

        Returns:
            torch.Tensor: The rotated tensor b_rot or (if b_only=False) a, b_rot, c.
        """
        K, B, _ = a.shape

        # Expand theta for compatibility with K and B dimensions
        theta = theta.unsqueeze(-1)  # Now shape (K, B, 1) for compatibility with broadcasting

        # Calculate normalized ac direction
        ac = c - a
        ac_norm = ac / torch.norm(ac, dim=2, keepdim=True)

        zeros = torch.zeros_like(theta)
        k = ac_norm

        K_matrix = torch.cat(
            [
                zeros,
                -k[..., 2:3],
                k[..., 1:2],
                k[..., 2:3],
                zeros,
                -k[..., 0:1],
                -k[..., 1:2],
                k[..., 0:1],
                zeros,
            ],
            dim=2,
        ).reshape(
            K * B, 3, 3
        )  # Reshape for batch matrix multiplication

        # Calculate rotation matrix R and the three terms
        R_1 = torch.eye(3, device=k.device).repeat(K * B, 1, 1)  # Repeat eye for each batch
        R_2 = torch.sin(theta).reshape(K * B, 1, 1) * K_matrix
        R_3 = (1 - torch.cos(theta).reshape(K * B, 1, 1)) * torch.bmm(K_matrix, K_matrix)
        R = R_1 + R_2 + R_3

        # Translate b to the origin (relative to a), then rotate, then translate back
        b_translated = (b - a).reshape(K * B, 3, 1)  # Reshape for batch matrix multiplication
        b_rot = torch.bmm(R, b_translated).reshape(K, B, 3) + a
        if b_only:
            return b_rot
        return a, b_rot, c
//  [python]
def rotate_ac(self, a, b, c, theta, b_only: bool = True):
        """Performs batch rotation for the tensor, withholding K x B x 3 tensors.

        Args:
            a, b, c (torch.Tensor): Input tensors of dimensions K x B x 3, representing initial position vectors for each batch.
            theta (torch.Tensor): The angle of rotation in radians. Vector based, size N x 1.
            b_only (bool): If True, only returns tensor b_rot, otherwise also returns a and c tensors. Default is True.

        Returns:
            torch.Tensor: The rotated tensor b_rot or (if b_only=False) a, b_rot, c.
        """
        K, B, _ = a.shape

        # Expand theta for compatibility with K and B dimensions
        theta = theta.unsqueeze(-1)  # Now shape (K, B, 1) for compatibility with broadcasting

        # Calculate normalized ac direction
        ac = c - a
        ac_norm = ac / torch.linalg.norm(ac, dim=2, keepdim=True)

        zeros = torch.zeros_like(theta)
        k = ac_norm

        K_matrix = torch.cat(
            [
                zeros,
                -k[..., 2:3],
                k[..., 1:2],
                k[..., 2:3],
                zeros,
                -k[..., 0:1],
                -k[..., 1:2],
                k[..., 0:1],
                zeros,
            ],
            dim=2,
        ).reshape(
            K * B, 3, 3
        )  # Reshape for batch matrix multiplication

        # Calculate rotation matrix R and the three terms
        R_1 = torch.eye(3, device=k.device).repeat(K * B, 1, 1)  # Repeat eye for each batch
        R_2 = torch.sin(theta).reshape(K * B, 1, 1) * K_matrix
        R_3 = (1 - torch.cos(theta).reshape(K * B, 1, 1)) * torch.bmm(K_matrix, K_matrix)
        R = R_1 + R_2 + R_3

        # Translate b to the origin (relative to a), then rotate, then translate back
        b_translated = (b - a).reshape(K * B, 3, 1)  # Reshape for batch matrix multiplication
        b_rot = torch.bmm(R, b_translated).reshape(K, B, 3) + a
        if b_only:
            return b_rot
        return a, b_rot, c
//  [python]
def forward(self, *args, **kwargs):
        outputs = self.model.forward(*args, **kwargs)

        if (getattr(outputs, "loss", None) is not None) and isinstance(outputs.loss, torch.Tensor):
            # Calculate the orthogonal regularization
            orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight

            if orth_reg_weight <= 0:
                raise ValueError("orth_reg_weight should be greater than 0. ")

            regu_loss = 0
            num_param = 0
            for n, p in self.model.named_parameters():
                if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n:
                    if p.shape == torch.Size([0]):
                        with gather_params_ctx(p, fwd_module=self):
                            para_cov = p @ p.T if "lora_A" in n else p.T @ p
                    else:
                        para_cov = p @ p.T if "lora_A" in n else p.T @ p
                    I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov))  # noqa: E741
                    I.requires_grad = False
                    num_param += 1
                    regu_loss += torch.norm(para_cov - I, p="fro")
            if num_param > 0:
                regu_loss = regu_loss / num_param
            else:
                regu_loss = 0
            outputs.loss += orth_reg_weight * regu_loss
        return outputs
//  [python]
def forward(self, *args, **kwargs):
        outputs = self.model.forward(*args, **kwargs)

        if (getattr(outputs, "loss", None) is not None) and isinstance(outputs.loss, torch.Tensor):
            # Calculate the orthogonal regularization
            orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight

            if orth_reg_weight <= 0:
                raise ValueError("orth_reg_weight should be greater than 0. ")

            regu_loss = 0
            num_param = 0
            for n, p in self.model.named_parameters():
                if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n:
                    if p.shape == torch.Size([0]):
                        with gather_params_ctx(p, fwd_module=self):
                            para_cov = p @ p.T if "lora_A" in n else p.T @ p
                    else:
                        para_cov = p @ p.T if "lora_A" in n else p.T @ p
                    I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov))  # noqa: E741
                    I.requires_grad = False
                    num_param += 1
                    regu_loss += torch.linalg.norm(para_cov - I, ord="fro")
            if num_param > 0:
                regu_loss = regu_loss / num_param
            else:
                regu_loss = 0
            outputs.loss += orth_reg_weight * regu_loss
        return outputs
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def ReprojectionError(p2d, p3d, P_13d, K, W):
    b = p2d.size(0)
    pts2d_pro = geo.project_points_by_theta(p3d, P_13d, K)
    if W is not None:
        error = torch.norm((W * (pts2d_pro - p2d)).reshape(b, -1), p=2, dim=-1) ** 2
    else:
        error = torch.norm(pts2d_pro.reshape(b, -1) - p2d.reshape(b, -1), p=2, dim=-1) ** 2
    return error
//  [python]
def ReprojectionError(p2d, p3d, P_13d, K, W):
    b = p2d.size(0)
    pts2d_pro = geo.project_points_by_theta(p3d, P_13d, K)
    if W is not None:
        error = torch.linalg.norm((W * (pts2d_pro - p2d)).reshape(b, -1), ord=2, dim=-1) ** 2
    else:
        error = torch.linalg.norm(pts2d_pro.reshape(b, -1) - p2d.reshape(b, -1), ord=2, dim=-1) ** 2
    return error
//  [python]
def ReprojectionError_flat(p2d, p3d, P_13d, K, W):
    b = p2d.size(0)
    pts2d_pro = geo.project_points_by_theta(p3d.view(b, -1, 3), P_13d, K)
    if W is not None:
        error = torch.norm((W.unsqueeze(-1) * (pts2d_pro - p2d.view(b, -1, 2))).reshape(b, -1), p=2, dim=-1) ** 2
    else:
        error = torch.norm(pts2d_pro.reshape(b, -1) - p2d, p=2, dim=-1) ** 2
    return error
//  [python]
def ReprojectionError_flat(p2d, p3d, P_13d, K, W):
    b = p2d.size(0)
    pts2d_pro = geo.project_points_by_theta(p3d.view(b, -1, 3), P_13d, K)
    if W is not None:
        error = torch.linalg.norm((W.unsqueeze(-1) * (pts2d_pro - p2d.view(b, -1, 2))).reshape(b, -1), ord=2, dim=-1) ** 2
    else:
        error = torch.linalg.norm(pts2d_pro.reshape(b, -1) - p2d, ord=2, dim=-1) ** 2
    return error
//  [python]
def L2norm(x):
    print(x.shape)
    return x / torch.norm(x, p=2, dim=2)
//  [python]
def L2norm(x):
    print(x.shape)
    return x / torch.linalg.norm(x, ord=2, dim=2)
//  [python]
def se3_compare(a, b):
    combined = se3_cat(a, se3_inv(b))

    trace = combined[..., 0, 0] + combined[..., 1, 1] + combined[..., 2, 2]
    rot_err_deg = torch.acos(torch.clamp(0.5 * (trace - 1), -1., 1.)) \
                  * 180 / math.pi
    trans_err = torch.norm(combined[..., :, 3], dim=-1)

    err = {
        'rot_deg': rot_err_deg,
        'trans': trans_err
    }
    return err
//  [python]
def se3_compare(a, b):
    combined = se3_cat(a, se3_inv(b))

    trace = combined[..., 0, 0] + combined[..., 1, 1] + combined[..., 2, 2]
    rot_err_deg = torch.acos(torch.clamp(0.5 * (trace - 1), -1., 1.)) \
                  * 180 / math.pi
    trans_err = torch.linalg.norm(combined[..., :, 3], dim=-1)

    err = {
        'rot_deg': rot_err_deg,
        'trans': trans_err
    }
    return err
//  [python]
def larc_layer_lr(p, lr, trust_coeff, wd, eps, clip=True, **kwargs):
    "Computes the local lr before weight decay is applied"
    p_norm,g_norm = torch.norm(p.data),torch.norm(p.grad.data)
    local_lr = lr*trust_coeff * (p_norm) / (g_norm + p_norm * wd + eps)
    return {'local_lr': min(lr, local_lr) if clip else local_lr}
//  [python]
def larc_layer_lr(p, lr, trust_coeff, wd, eps, clip=True, **kwargs):
    "Computes the local lr before weight decay is applied"
    p_norm,g_norm = torch.linalg.norm(p.data),torch.linalg.norm(p.grad.data)
    local_lr = lr*trust_coeff * (p_norm) / (g_norm + p_norm * wd + eps)
    return {'local_lr': min(lr, local_lr) if clip else local_lr}
//  [python]
def test_norm_old_nan_propagation(self, device):
        ords = [inf, -inf]
        for pair in itertools.product([0.0, nan, 1.0], repeat=2):
            x = torch.tensor(list(pair), device=device)
            for ord in ords:
                result = torch.norm(x, p=ord)
                result_check = torch.linalg.norm(x, ord=ord)
                self.assertEqual(result, result_check)
//  [python]
def test_norm_old_nan_propagation(self, device):
        ords = [inf, -inf]
        for pair in itertools.product([0.0, nan, 1.0], repeat=2):
            x = torch.tensor(list(pair), device=device)
            for ord in ords:
                result = torch.linalg.norm(x, ord=ord)
                result_check = torch.linalg.norm(x, ord=ord)
                self.assertEqual(result, result_check)
//  [python]
def forward(ctx, goal_quat, current_quat, quat_res):
        quat_res = OrientationError.geodesic_distance(goal_quat, current_quat, quat_res)
        rot_error = torch.norm(quat_res, dim=-1, keepdim=True)
        ctx.save_for_backward(quat_res, rot_error)
        return rot_error
//  [python]
def forward(ctx, goal_quat, current_quat, quat_res):
        quat_res = OrientationError.geodesic_distance(goal_quat, current_quat, quat_res)
        rot_error = torch.linalg.norm(quat_res, dim=-1, keepdim=True)
        ctx.save_for_backward(quat_res, rot_error)
        return rot_error
//  [python]
def test_fisher_regression_fd(self):
        model = torch.nn.Sequential(torch.nn.Linear(7500, 1, bias=False))

        data = torch.randn(50, 7500)

        # closed form fisher information for linear regression
        fim = data.t() @ data / data.size(0)

        eps_list = [1.0, 1e-1, 5e-2, 1e-2, 1e-3, 5e-3, 1e-4, 5e-4, 1e-5, 5e-6, 1e-6, 5e-6]

        for epsilon in eps_list:
            fvp = FVPR_FD(model, data, epsilon=epsilon)
            fvp_approx = fvp._matmul(torch.eye(data.shape[-1]))
            rel_error = torch.norm(fvp_approx - fim) / torch.norm(fim)

            # check to ensure that our relative error is never greater than 4%
            self.assertLess(rel_error.item(), 4e-2)
//  [python]
def test_fisher_regression_fd(self):
        model = torch.nn.Sequential(torch.nn.Linear(7500, 1, bias=False))

        data = torch.randn(50, 7500)

        # closed form fisher information for linear regression
        fim = data.t() @ data / data.size(0)

        eps_list = [1.0, 1e-1, 5e-2, 1e-2, 1e-3, 5e-3, 1e-4, 5e-4, 1e-5, 5e-6, 1e-6, 5e-6]

        for epsilon in eps_list:
            fvp = FVPR_FD(model, data, epsilon=epsilon)
            fvp_approx = fvp._matmul(torch.eye(data.shape[-1]))
            rel_error = torch.linalg.norm(fvp_approx - fim) / torch.linalg.norm(fim)

            # check to ensure that our relative error is never greater than 4%
            self.assertLess(rel_error.item(), 4e-2)
//  [python]
def frobenius_norm(self):
        return torch.norm(self.data)
//  [python]
def frobenius_norm(self):
        return torch.linalg.norm(self.data)
//  [python]
def grad_loss_term(depth_est, depth_gt, mask, stage=None):
    """
    Scale-invariant gradient loss as Eq.(12) in DeepTAM
    """
    grad_loss = 0.
    grad_deltas = [1, 2, 4]
    for h in grad_deltas:
        grad_gt = get_depth_grad_img(depth_gt, h).detach()
        grad_est = get_depth_grad_img(depth_est, h)
        # TODO: the mask of gradient shouldn't be the same as the original mask. The pixels left and down to the invalid points should not be counted as well.
        norm = (torch.norm(grad_est - grad_gt, dim=1)) * mask
        grad_loss += norm.mean() / mask.mean()
    grad_loss = grad_loss / len(grad_deltas)
    return grad_loss
//  [python]
def grad_loss_term(depth_est, depth_gt, mask, stage=None):
    """
    Scale-invariant gradient loss as Eq.(12) in DeepTAM
    """
    grad_loss = 0.
    grad_deltas = [1, 2, 4]
    for h in grad_deltas:
        grad_gt = get_depth_grad_img(depth_gt, h).detach()
        grad_est = get_depth_grad_img(depth_est, h)
        # TODO: the mask of gradient shouldn't be the same as the original mask. The pixels left and down to the invalid points should not be counted as well.
        norm = (torch.linalg.norm(grad_est - grad_gt, dim=1)) * mask
        grad_loss += norm.mean() / mask.mean()
    grad_loss = grad_loss / len(grad_deltas)
    return grad_loss
//  [python]
def frobenius_norm(self):
        # TODO test
        return sum([torch.norm(b)**2 for b in self.data.values()])**.5
//  [python]
def frobenius_norm(self):
        # TODO test
        return sum([torch.linalg.norm(b)**2 for b in self.data.values()])**.5
//  [python]
def frobenius_norm(self):
        A = torch.mm(self.data.view(-1, self.data.size(2)),
                     self.data.view(-1, self.data.size(2)).t())
        return torch.norm(A)
//  [python]
def frobenius_norm(self):
        A = torch.mm(self.data.view(-1, self.data.size(2)),
                     self.data.view(-1, self.data.size(2)).t())
        return torch.linalg.norm(A)
//  [python]
def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][
            0].device  # put everything on the same device, in case of model parallelism
        norm = torch.norm(
            torch.stack([
                ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for group in self.param_groups for p in group["params"]
                if p.grad is not None
            ]),
            p=2
        )
        return norm
//  [python]
def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][
            0].device  # put everything on the same device, in case of model parallelism
        norm = torch.linalg.norm(
            torch.stack([
                ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for group in self.param_groups for p in group["params"]
                if p.grad is not None
            ]),
            p=2
        )
        return norm
//  [python]
def Slerp(v0, v1, t):
    unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
    v0 = unit_vec(v0)
    v1 = unit_vec(v1)
    dot = torch.sum(v0 * v1, dim=-1)
    
    # Calculate initial angle between v0 and v1
    omega = torch.arccos(dot)
    sin_omega = torch.sin(omega)
    
    # Angle at timestep t
    s0 = torch.sin((1-t)*omega) / sin_omega
    s1 = torch.sin(t*omega) / sin_omega
    v2 = s0[...,None] * v0 + s1[...,None] * v1
    return v2
//  [python]
def Slerp(v0, v1, t):
    unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
    v0 = unit_vec(v0)
    v1 = unit_vec(v1)
    dot = torch.sum(v0 * v1, dim=-1)
    
    # Calculate initial angle between v0 and v1
    omega = torch.arccos(dot)
    sin_omega = torch.sin(omega)
    
    # Angle at timestep t
    s0 = torch.sin((1-t)*omega) / sin_omega
    s1 = torch.sin(t*omega) / sin_omega
    v2 = s0[...,None] * v0 + s1[...,None] * v1
    return v2
//  [python]
def angle(self, v1, v2):
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        v1 = unit_vec(v1)
        v2 = unit_vec(v2)
        cos = torch.sum(v1*v2, dim=-1)
        cos = torch.clamp(cos, -0.99999,0.99999)
        return torch.arccos(cos)
//  [python]
def angle(self, v1, v2):
        unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
        v1 = unit_vec(v1)
        v2 = unit_vec(v2)
        cos = torch.sum(v1*v2, dim=-1)
        cos = torch.clamp(cos, -0.99999,0.99999)
        return torch.arccos(cos)
//  [python]
def dihedral(self, v1, v2, v3):
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        v1 = unit_vec(v1)
        v2 = unit_vec(v2)
        v3 = unit_vec(v3)
        
        n1 = torch.cross(v1, v2, dim=-1)
        n2 = torch.cross(v2, v3, dim=-1)
        x = torch.sum(n1*n2, dim=-1)
        y = torch.sum(torch.cross(n1,n2, dim=-1)*v2, dim=-1)
        return torch.arctan2(y,x)
//  [python]
def dihedral(self, v1, v2, v3):
        unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
        v1 = unit_vec(v1)
        v2 = unit_vec(v2)
        v3 = unit_vec(v3)
        
        n1 = torch.cross(v1, v2, dim=-1)
        n2 = torch.cross(v2, v3, dim=-1)
        x = torch.sum(n1*n2, dim=-1)
        y = torch.sum(torch.cross(n1,n2, dim=-1)*v2, dim=-1)
        return torch.arctan2(y,x)
//  [python]
def sqrtm_torch_acc(m):

    assert torch.allclose(m, m.t()), "The matrix is not symmetric."
    e, v = m.symeig(eigenvectors=True)
    root = torch.mm(torch.mm(v, e.sqrt().diag()), v.t())
    # embed()
    print('\tMatrix square root: distance {:2.6f}, matrix m norm {:2.6f}'.
          format(torch.dist(m, torch.mm(root, root)).cpu().detach().numpy(), tol * torch.norm(m).cpu().detach().numpy()))
    assert torch.dist(m, torch.mm(root, root)) <= (tol * torch.norm(m)), "The result is not the matrix square root."
    return root
//  [python]
def sqrtm_torch_acc(m):

    assert torch.allclose(m, m.t()), "The matrix is not symmetric."
    e, v = m.symeig(eigenvectors=True)
    root = torch.mm(torch.mm(v, e.sqrt().diag()), v.t())
    # embed()
    print('\tMatrix square root: distance {:2.6f}, matrix m norm {:2.6f}'.
          format(torch.dist(m, torch.mm(root, root)).cpu().detach().numpy(), tol * torch.linalg.norm(m).cpu().detach().numpy()))
    assert torch.dist(m, torch.mm(root, root)) <= (tol * torch.linalg.norm(m)), "The result is not the matrix square root."
    return root
//  [python]
def sqrtm_torch(m):

    assert torch.allclose(m, m.t()), "The matrix is not symmetric."
    u, s, v = m.svd()
    root = torch.mm(torch.mm(u, s.sqrt().diag()), v.t())
    print('\tMatrix square root: distance {:2.6f}, matrix m norm {:2.6f}'.
          format(torch.dist(m, torch.mm(root, root)).cpu().detach().numpy(), tol * torch.norm(m).cpu().detach().numpy()))
    assert torch.dist(m, torch.mm(root, root)) <= (tol * torch.norm(m)), "The result is not the matrix square root."
    return root
//  [python]
def sqrtm_torch(m):

    assert torch.allclose(m, m.t()), "The matrix is not symmetric."
    u, s, v = m.svd()
    root = torch.mm(torch.mm(u, s.sqrt().diag()), v.t())
    print('\tMatrix square root: distance {:2.6f}, matrix m norm {:2.6f}'.
          format(torch.dist(m, torch.mm(root, root)).cpu().detach().numpy(), tol * torch.linalg.norm(m).cpu().detach().numpy()))
    assert torch.dist(m, torch.mm(root, root)) <= (tol * torch.linalg.norm(m)), "The result is not the matrix square root."
    return root