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

        min_cost = np.inf

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

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

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

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

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

        return (center.unsqueeze(0) * (x_max - x_min) + x_min)
//  [python]
def normalize(self, x):
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 0.00001)
        return x_normalized
//  [python]
def normalize(self, x):
        x_norm = torch.linalg.norm(x, ord=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 0.00001)
        return x_normalized
//  [python]
def forward(self, x, enc_len=None, returnClusterIds=False):
        if self.reestimationReservoir and self.training and x.device.index == 0:
            self.reestimate()

        if self.logInputNorms and self.training:
            norms = torch.norm(x.contiguous().view(-1, x.size(-1)), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms pre-BN/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms pre-BN/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        
        bs, l, d = x.size()
        x = x.transpose(1, 2)
        x = self.batchNorm(x)
        x = x.transpose(1, 2).contiguous()

        if self.logInputNorms and self.training:
            norms = torch.norm(x.view(-1, x.size(-1)), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms post-BN/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms post-BN/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        # x = self.projection(x)
        if self.training and self.reestimationReservoir and x.device.index == 0:
            self.reestimationReservoir.add(
                self.pack_x(x, enc_len)
                    .view(-1, x.size(-1)).detach())

        if Globals.epoch < self.quantizationEnforceFromEpoch and not returnClusterIds:
            # print("Skipping quantization")
            codes = x
            indices = torch.zeros(x.shape[:-1] + (1,), device=x.device, dtype=torch.int64)
            values = indices
        else:
            x = self.pack_x(x, enc_len)
            codes, indices = quantize(x, self.embedding.weight, self.commitment, 'nearest', False)
            codes = self.unpack_x(codes, enc_len)
            indices = self.unpack_x(indices, enc_len)

        if self.training:
            self._logCodeUsage(indices, 'nearest')

        if self.logInputNorms and self.training:
            norms = torch.norm(codes.view(-1, codes.size(-1)).contiguous(), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms post-VQ/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms post-VQ/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        if returnClusterIds:
            return indices
        return codes
//  [python]
def forward(self, x, enc_len=None, returnClusterIds=False):
        if self.reestimationReservoir and self.training and x.device.index == 0:
            self.reestimate()

        if self.logInputNorms and self.training:
            norms = torch.linalg.norm(x.contiguous().view(-1, x.size(-1)), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms pre-BN/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms pre-BN/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        
        bs, l, d = x.size()
        x = x.transpose(1, 2)
        x = self.batchNorm(x)
        x = x.transpose(1, 2).contiguous()

        if self.logInputNorms and self.training:
            norms = torch.linalg.norm(x.view(-1, x.size(-1)), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms post-BN/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms post-BN/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        # x = self.projection(x)
        if self.training and self.reestimationReservoir and x.device.index == 0:
            self.reestimationReservoir.add(
                self.pack_x(x, enc_len)
                    .view(-1, x.size(-1)).detach())

        if Globals.epoch < self.quantizationEnforceFromEpoch and not returnClusterIds:
            # print("Skipping quantization")
            codes = x
            indices = torch.zeros(x.shape[:-1] + (1,), device=x.device, dtype=torch.int64)
            values = indices
        else:
            x = self.pack_x(x, enc_len)
            codes, indices = quantize(x, self.embedding.weight, self.commitment, 'nearest', False)
            codes = self.unpack_x(codes, enc_len)
            indices = self.unpack_x(indices, enc_len)

        if self.training:
            self._logCodeUsage(indices, 'nearest')

        if self.logInputNorms and self.training:
            norms = torch.linalg.norm(codes.view(-1, codes.size(-1)).contiguous(), dim=1)
            Globals.writer.add_scalar(tag='VQ Input norms post-VQ/mean', scalar_value=torch.mean(norms), global_step=Globals.currentIteration)
            Globals.writer.add_scalar(tag='VQ Input norms post-VQ/std', scalar_value=torch.std(norms), global_step=Globals.currentIteration)

        if returnClusterIds:
            return indices
        return codes
//  [python]
def plot_pca(self, ax, manifold_collection, title=None, point_color=None, draw_line=True, draw_arrow=True):
        ax.cla()

        point_alpha = 0.3
        line_alpha = 0.2
        arrow_alpha = 1.0
        arrow_step = 50
        arrow_size = 0.015
        arrow_power = 1.0
        arrow_color = (0.25, 0.25, 0.5)

        num_steps = [len(manifold) for manifold in manifold_collection]
        manifolds = torch.cat(manifold_collection, dim=0).cpu()
        
        manifolds_pca = torch.tensor(self.pca.fit_transform(manifolds)).split(num_steps, dim=0)

        for i, manifold in enumerate(manifolds_pca):
            if draw_line:
                ax.plot(manifold[:, 0], manifold[:, 1], color=(0.0, 0.0, 0.0), alpha=line_alpha)
            if point_color is None:
                ax.scatter(manifold[:, 0], manifold[:, 1], alpha=point_alpha, label=i)
            else:
                ax.scatter(manifold[:, 0], manifold[:, 1], color=point_color[i], alpha=point_alpha, label=i)
            if draw_arrow:
                for j in range(0, len(manifold) - arrow_step, arrow_step):
                    d = torch.norm(manifold[j, :], dim=-1)
                    d = torch.pow(d, arrow_power)
                    ax.arrow(
                        manifold[j, 0],
                        manifold[j, 1],
                        manifold[j + 1, 0] - manifold[j, 0],
                        manifold[j + 1, 1] - manifold[j, 1],
                        alpha=arrow_alpha, width=d * arrow_size, color=arrow_color
                        )
        ax.legend()
        ax.set_axis_off()
        if title != None:
            ax.set_title(title)
//  [python]
def plot_pca(self, ax, manifold_collection, title=None, point_color=None, draw_line=True, draw_arrow=True):
        ax.cla()

        point_alpha = 0.3
        line_alpha = 0.2
        arrow_alpha = 1.0
        arrow_step = 50
        arrow_size = 0.015
        arrow_power = 1.0
        arrow_color = (0.25, 0.25, 0.5)

        num_steps = [len(manifold) for manifold in manifold_collection]
        manifolds = torch.cat(manifold_collection, dim=0).cpu()
        
        manifolds_pca = torch.tensor(self.pca.fit_transform(manifolds)).split(num_steps, dim=0)

        for i, manifold in enumerate(manifolds_pca):
            if draw_line:
                ax.plot(manifold[:, 0], manifold[:, 1], color=(0.0, 0.0, 0.0), alpha=line_alpha)
            if point_color is None:
                ax.scatter(manifold[:, 0], manifold[:, 1], alpha=point_alpha, label=i)
            else:
                ax.scatter(manifold[:, 0], manifold[:, 1], color=point_color[i], alpha=point_alpha, label=i)
            if draw_arrow:
                for j in range(0, len(manifold) - arrow_step, arrow_step):
                    d = torch.linalg.norm(manifold[j, :], dim=-1)
                    d = torch.pow(d, arrow_power)
                    ax.arrow(
                        manifold[j, 0],
                        manifold[j, 1],
                        manifold[j + 1, 0] - manifold[j, 0],
                        manifold[j + 1, 1] - manifold[j, 1],
                        alpha=arrow_alpha, width=d * arrow_size, color=arrow_color
                        )
        ax.legend()
        ax.set_axis_off()
        if title != None:
            ax.set_title(title)
//  [python]
def compute_world2aligned_joints_mat(joints):
    """
    Compute world to canonical frame (rotation around up axis)
    from the given batch of joints (B x J x 3)
    """
    left_idx = SMPL_JOINTS["leftUpLeg"]
    right_idx = SMPL_JOINTS["rightUpLeg"]

    body_right = joints[:, right_idx] - joints[:, left_idx]
    body_right = body_right / torch.norm(body_right, dim=1, keepdim=True)

    world2aligned_mat, world2aligned_aa = compute_aligned_from_right(body_right)

    return world2aligned_mat
//  [python]
def compute_world2aligned_joints_mat(joints):
    """
    Compute world to canonical frame (rotation around up axis)
    from the given batch of joints (B x J x 3)
    """
    left_idx = SMPL_JOINTS["leftUpLeg"]
    right_idx = SMPL_JOINTS["rightUpLeg"]

    body_right = joints[:, right_idx] - joints[:, left_idx]
    body_right = body_right / torch.linalg.norm(body_right, dim=1, keepdim=True)

    world2aligned_mat, world2aligned_aa = compute_aligned_from_right(body_right)

    return world2aligned_mat
//  [python]
def forward(self, x):
        def norm(t):
            n = torch.norm(t, dim=-1, keepdim=True).clamp(min=self.eps)
            return t / n * self.g
        return map_first_tuple_or_el(x, norm)
//  [python]
def forward(self, x):
        def norm(t):
            n = torch.linalg.norm(t, dim=-1, keepdim=True).clamp(min=self.eps)
            return t / n * self.g
        return map_first_tuple_or_el(x, norm)
//  [python]
def compute_aligned_from_right(body_right):
    xy_axis = XY_AXIS_GLOB
    x_axis = X_AXIS_GLOB

    body_right_x_proj = body_right[:, 0:1] / (
        torch.norm(body_right[:, :2], dim=1, keepdim=True) + 1e-6
    )
    body_right_x_proj = torch.clamp(
        body_right_x_proj, min=-1.0, max=1.0
    )  # avoid acos error

    world2aligned_angle = torch.acos(
        body_right_x_proj
    )  # project to world x axis, and compute angle
    body_right = body_right * xy_axis
    world2aligned_axis = torch.linalg.cross(body_right, x_axis.expand_as(body_right))

    world2aligned_aa = (
        world2aligned_axis
        / (torch.norm(world2aligned_axis, dim=1, keepdim=True) + 1e-6)
    ) * world2aligned_angle
    world2aligned_mat = batch_rodrigues(world2aligned_aa)

    return world2aligned_mat, world2aligned_aa
//  [python]
def compute_aligned_from_right(body_right):
    xy_axis = XY_AXIS_GLOB
    x_axis = X_AXIS_GLOB

    body_right_x_proj = body_right[:, 0:1] / (
        torch.linalg.norm(body_right[:, :2], dim=1, keepdim=True) + 1e-6
    )
    body_right_x_proj = torch.clamp(
        body_right_x_proj, min=-1.0, max=1.0
    )  # avoid acos error

    world2aligned_angle = torch.acos(
        body_right_x_proj
    )  # project to world x axis, and compute angle
    body_right = body_right * xy_axis
    world2aligned_axis = torch.linalg.cross(body_right, x_axis.expand_as(body_right))

    world2aligned_aa = (
        world2aligned_axis
        / (torch.linalg.norm(world2aligned_axis, dim=1, keepdim=True) + 1e-6)
    ) * world2aligned_angle
    world2aligned_mat = batch_rodrigues(world2aligned_aa)

    return world2aligned_mat, world2aligned_aa
//  [python]
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    """Calculates the rotation matrices for a batch of rotation vectors
    Parameters
    ----------
    rot_vecs: torch.tensor Nx3
        array of N axis-angle vectors
    Returns
    -------
    R: torch.tensor Nx3x3
        The rotation matrices for the given axis-angle parameters
    """

    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):
    """Calculates the rotation matrices for a batch of rotation vectors
    Parameters
    ----------
    rot_vecs: torch.tensor Nx3
        array of N axis-angle vectors
    Returns
    -------
    R: torch.tensor Nx3x3
        The rotation matrices for the given axis-angle parameters
    """

    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 solve_multilayer_translation(query, positive, rot_angle, device):
    B, H, W = query.shape
    # caculate the translation of b relative to a
    angles = torch.FloatTensor(np.linspace(0, 2*np.pi, H).astype(np.float32)).to(device)
    angles = angles + rot_angle
    # matrices of the overdetermined linear system
    A = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
    b = torch.FloatTensor(H).to(device)
    # calculate the condition number of A
    cond = torch.linalg.cond(A)
    print("cond: ", cond)

    for i in range(H):
        query_fft = torch.fft.fft2(query[:,i,:], dim=-1, norm="ortho")
        positive_fft = torch.fft.fft2(positive[:,i,:], dim=-1, norm="ortho")
        corr = torch.fft.ifft2(query_fft*positive_fft.conj(), dim=-1, norm="ortho")
        corr = torch.sqrt(corr.imag**2 + corr.real**2)
        corr = fftshift1d(corr)
        corr = torch.sum(corr,dim=0)
        shift = torch.argmax(corr) - W//2
        b[i] = shift

    x, y = solve_overdetermined_linear_system(A, b, method='svd')
    # overdetermined linear system error
    error = torch.norm(torch.matmul(A, torch.cat([x, y], dim=0)) - b)
    x = x.cpu().numpy()
    y = y.cpu().numpy()
    error = error.cpu().numpy()
    
    return x, y, error
//  [python]
def solve_multilayer_translation(query, positive, rot_angle, device):
    B, H, W = query.shape
    # caculate the translation of b relative to a
    angles = torch.FloatTensor(np.linspace(0, 2*np.pi, H).astype(np.float32)).to(device)
    angles = angles + rot_angle
    # matrices of the overdetermined linear system
    A = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
    b = torch.FloatTensor(H).to(device)
    # calculate the condition number of A
    cond = torch.linalg.cond(A)
    print("cond: ", cond)

    for i in range(H):
        query_fft = torch.fft.fft2(query[:,i,:], dim=-1, norm="ortho")
        positive_fft = torch.fft.fft2(positive[:,i,:], dim=-1, norm="ortho")
        corr = torch.fft.ifft2(query_fft*positive_fft.conj(), dim=-1, norm="ortho")
        corr = torch.sqrt(corr.imag**2 + corr.real**2)
        corr = fftshift1d(corr)
        corr = torch.sum(corr,dim=0)
        shift = torch.argmax(corr) - W//2
        b[i] = shift

    x, y = solve_overdetermined_linear_system(A, b, method='svd')
    # overdetermined linear system error
    error = torch.linalg.norm(torch.matmul(A, torch.cat([x, y], dim=0)) - b)
    x = x.cpu().numpy()
    y = y.cpu().numpy()
    error = error.cpu().numpy()
    
    return x, y, error
//  [python]
def solve_translation(query, positive, rot_angle, device):
    B, H, W = query.shape
    # caculate the translation of b relative to a
    angles = torch.FloatTensor(np.linspace(0, 2*np.pi, H).astype(np.float32)).to(device)
    # compensate for the rotation of the query to calculate the translation
    angles = angles + rot_angle # in radians
    # # take one half of the spectrum for translation estimation
    # angles = angles[0:H//2]
    # query = query[0:H//2,:]
    # positive = positive[0:H//2,:]
    # matrices of the overdetermined linear system
    A = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
    b = torch.FloatTensor(H).to(device)

    # calculate the condition number of A
    cond = torch.linalg.cond(A)

    for i in range(H):
        query_fft = torch.fft.fft2(query[:,i,:], dim=-1, norm="ortho")
        positive_fft = torch.fft.fft2(positive[:,i,:], dim=-1, norm="ortho")
        corr = torch.fft.ifft2(query_fft*positive_fft.conj(), dim=-1, norm="ortho")
        corr = torch.sqrt(corr.imag**2 + corr.real**2)
        # use the torch fftshift function to shift the correlation to the center (much faster)
        corr = torch.fft.fftshift(corr)  
        corr = torch.sum(corr,dim=0)
        shift = W//2 - torch.argmax(corr)
        b[i] = shift
    x, y = solve_overdetermined_linear_system(A, b, method='svd')
    # overdetermined linear system error
    error = torch.norm(torch.matmul(A, torch.cat([x, y], dim=0)) - b)
    x = x.cpu().numpy()
    y = y.cpu().numpy()
    error = error.cpu().numpy()
    # print('predicted x, y, error: ', x, y, error)
    
    return x, y, error
//  [python]
def solve_translation(query, positive, rot_angle, device):
    B, H, W = query.shape
    # caculate the translation of b relative to a
    angles = torch.FloatTensor(np.linspace(0, 2*np.pi, H).astype(np.float32)).to(device)
    # compensate for the rotation of the query to calculate the translation
    angles = angles + rot_angle # in radians
    # # take one half of the spectrum for translation estimation
    # angles = angles[0:H//2]
    # query = query[0:H//2,:]
    # positive = positive[0:H//2,:]
    # matrices of the overdetermined linear system
    A = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
    b = torch.FloatTensor(H).to(device)

    # calculate the condition number of A
    cond = torch.linalg.cond(A)

    for i in range(H):
        query_fft = torch.fft.fft2(query[:,i,:], dim=-1, norm="ortho")
        positive_fft = torch.fft.fft2(positive[:,i,:], dim=-1, norm="ortho")
        corr = torch.fft.ifft2(query_fft*positive_fft.conj(), dim=-1, norm="ortho")
        corr = torch.sqrt(corr.imag**2 + corr.real**2)
        # use the torch fftshift function to shift the correlation to the center (much faster)
        corr = torch.fft.fftshift(corr)  
        corr = torch.sum(corr,dim=0)
        shift = W//2 - torch.argmax(corr)
        b[i] = shift
    x, y = solve_overdetermined_linear_system(A, b, method='svd')
    # overdetermined linear system error
    error = torch.linalg.norm(torch.matmul(A, torch.cat([x, y], dim=0)) - b)
    x = x.cpu().numpy()
    y = y.cpu().numpy()
    error = error.cpu().numpy()
    # print('predicted x, y, error: ', x, y, error)
    
    return x, y, error
//  [python]
def norm(
        self,
        p: Optional[Union[float, str]] = "fro",
        dim=None,
        keepdim=False,
        dtype=None,
    ):
        r"""See :func:`torch.norm`"""
        if has_torch_function_unary(self):
            return handle_torch_function(
                Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype
            )
        return torch.norm(self, p, dim, keepdim, dtype=dtype)
//  [python]
def norm(
        self,
        p: Optional[Union[float, str]] = "fro",
        dim=None,
        keepdim=False,
        dtype=None,
    ):
        r"""See :func:`torch.linalg.norm`"""
        if has_torch_function_unary(self):
            return handle_torch_function(
                Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype
            )
        return torch.linalg.norm(self, p, dim, keepdim, dtype=dtype)
//  [python]
def kernel_normalize(self, kernel, k=3):
        b, ch, h, w, kk = kernel.size()
        kernel = kernel.view(b, ch, h * w, kk).transpose(2, 1)
        kernel_norm = torch.norm(kernel.contiguous().view(b, h * w, ch * kk), p=2, dim=2, keepdim=True)
        kernel = kernel.view(b, h * w, ch, k, k)
        kernel_norm = kernel_norm.view(b, h * w, 1, 1, 1)
        return kernel, kernel_norm
//  [python]
def kernel_normalize(self, kernel, k=3):
        b, ch, h, w, kk = kernel.size()
        kernel = kernel.view(b, ch, h * w, kk).transpose(2, 1)
        kernel_norm = torch.linalg.norm(kernel.contiguous().view(b, h * w, ch * kk), ord=2, dim=2, keepdim=True)
        kernel = kernel.view(b, h * w, ch, k, k)
        kernel_norm = kernel_norm.view(b, h * w, 1, 1, 1)
        return kernel, kernel_norm
//  [python]
def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x
//  [python]
def forward(self, x):
        Gx = torch.linalg.norm(x, ord=2, dim=(2, 3), keepdim=True)
        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x
//  [python]
def test_jacobian_kfac():
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()

        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        M_kfac = PMatKFAC(generator=generator, examples=loader)
        G_kfac_split = M_kfac.get_dense_tensor(split_weight_bias=True)
        G_kfac = M_kfac.get_dense_tensor(split_weight_bias=False)

        # Test trace
        trace_direct = torch.trace(G_kfac_split)
        trace_kfac = M_kfac.trace()
        check_ratio(trace_direct, trace_kfac)

        # Test frobenius norm
        frob_direct = torch.norm(G_kfac)
        frob_kfac = M_kfac.frobenius_norm()
        check_ratio(frob_direct, frob_kfac)

        # Test get_diag
        check_tensors(torch.diag(G_kfac_split), M_kfac.get_diag(split_weight_bias=True))

        # sample random vector
        random_v = random_pvector(lc, device)

        # Test mv
        mv_direct = torch.mv(G_kfac_split, random_v.get_flat_representation())
        mv_kfac = M_kfac.mv(random_v)
        check_tensors(mv_direct, mv_kfac.get_flat_representation())

        # Test vTMv
        mnorm_kfac = M_kfac.vTMv(random_v)
        mnorm_direct = torch.dot(mv_direct, random_v.get_flat_representation())
        check_ratio(mnorm_direct, mnorm_kfac)

        # Test inverse
        # We start from a mv vector since it kills its components projected to
        # the small eigenvalues of KFAC
        regul = 1e-7

        mv2 = M_kfac.mv(mv_kfac)
        kfac_inverse = M_kfac.inverse(regul)
        mv_back = kfac_inverse.mv(mv2 + regul * mv_kfac)
        check_tensors(
            mv_kfac.get_flat_representation(),
            mv_back.get_flat_representation(),
            eps=1e-2,
        )

        # Test solve
        mv_back = M_kfac.solve(mv2 + regul * mv_kfac, regul=regul)
        check_tensors(
            mv_kfac.get_flat_representation(),
            mv_back.get_flat_representation(),
            eps=1e-2,
        )
//  [python]
def test_jacobian_kfac():
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()

        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        M_kfac = PMatKFAC(generator=generator, examples=loader)
        G_kfac_split = M_kfac.get_dense_tensor(split_weight_bias=True)
        G_kfac = M_kfac.get_dense_tensor(split_weight_bias=False)

        # Test trace
        trace_direct = torch.trace(G_kfac_split)
        trace_kfac = M_kfac.trace()
        check_ratio(trace_direct, trace_kfac)

        # Test frobenius norm
        frob_direct = torch.linalg.norm(G_kfac)
        frob_kfac = M_kfac.frobenius_norm()
        check_ratio(frob_direct, frob_kfac)

        # Test get_diag
        check_tensors(torch.diag(G_kfac_split), M_kfac.get_diag(split_weight_bias=True))

        # sample random vector
        random_v = random_pvector(lc, device)

        # Test mv
        mv_direct = torch.mv(G_kfac_split, random_v.get_flat_representation())
        mv_kfac = M_kfac.mv(random_v)
        check_tensors(mv_direct, mv_kfac.get_flat_representation())

        # Test vTMv
        mnorm_kfac = M_kfac.vTMv(random_v)
        mnorm_direct = torch.dot(mv_direct, random_v.get_flat_representation())
        check_ratio(mnorm_direct, mnorm_kfac)

        # Test inverse
        # We start from a mv vector since it kills its components projected to
        # the small eigenvalues of KFAC
        regul = 1e-7

        mv2 = M_kfac.mv(mv_kfac)
        kfac_inverse = M_kfac.inverse(regul)
        mv_back = kfac_inverse.mv(mv2 + regul * mv_kfac)
        check_tensors(
            mv_kfac.get_flat_representation(),
            mv_back.get_flat_representation(),
            eps=1e-2,
        )

        # Test solve
        mv_back = M_kfac.solve(mv2 + regul * mv_kfac, regul=regul)
        check_tensors(
            mv_kfac.get_flat_representation(),
            mv_back.get_flat_representation(),
            eps=1e-2,
        )
//  [python]
def test_norm_fastpaths(self):
        x = torch.randn(3, 5)

        # slow path
        result = torch.norm(x, 4.5, 1)
        expected = torch.pow(x.abs().pow(4.5).sum(1), 1.0 / 4.5)
        self.assertEqual(result, expected)

        # fast 0-norm
        result = torch.norm(x, 0, 1)
        expected = (x != 0).type_as(x).sum(1)
        self.assertEqual(result, expected)

        # fast 1-norm
        result = torch.norm(x, 1, 1)
        expected = x.abs().sum(1)
        self.assertEqual(result, expected)

        # fast 2-norm
        result = torch.norm(x, 2, 1)
        expected = torch.sqrt(x.pow(2).sum(1))
        self.assertEqual(result, expected)

        # fast 3-norm
        result = torch.norm(x, 3, 1)
        expected = torch.pow(x.pow(3).abs().sum(1), 1.0 / 3.0)
        self.assertEqual(result, expected)
//  [python]
def test_norm_fastpaths(self):
        x = torch.randn(3, 5)

        # slow path
        result = torch.linalg.norm(x, ord=4.5, dim=1)
        expected = torch.pow(x.abs().pow(4.5).sum(1), 1.0 / 4.5)
        self.assertEqual(result, expected)

        # fast 0-norm
        result = torch.linalg.norm(x, ord=0, dim=1)
        expected = (x != 0).type_as(x).sum(1)
        self.assertEqual(result, expected)

        # fast 1-norm
        result = torch.linalg.norm(x, ord=1, dim=1)
        expected = x.abs().sum(1)
        self.assertEqual(result, expected)

        # fast 2-norm
        result = torch.linalg.norm(x, ord=2, dim=1)
        expected = torch.sqrt(x.pow(2).sum(1))
        self.assertEqual(result, expected)

        # fast 3-norm
        result = torch.linalg.norm(x, ord=3, dim=1)
        expected = torch.pow(x.pow(3).abs().sum(1), 1.0 / 3.0)
        self.assertEqual(result, expected)
//  [python]
def area_loss(self, vert):
        N, V, nver, _ = vert.shape
        vert = vert * 1e-3

        # Compute area per view
        tri_mask = self.bfm.face_region_tri_mask.ravel()
        pt0 = vert[:, :, self.bfm.model['tri'][tri_mask, 2], :]                 # (N, V, ntri_masked, 3)
        pt1 = vert[:, :, self.bfm.model['tri'][tri_mask, 1], :]                 # (N, V, ntri_masked, 3)
        pt2 = vert[:, :, self.bfm.model['tri'][tri_mask, 0], :]                 # (N, V, ntri_masked, 3)
        tri_area = torch.norm(torch.cross(pt0 - pt1, pt0 - pt2, dim=3), dim=3)  # (N, V, ntri_masked)
        # tri_area_max, _ = torch.max(tri_area, dim=1, keepdim=True)
        # tri_area_max = tri_area_max.detach()                                    # (N, 1, ntri_masked)
        #
        # # Compute loss
        # loss = torch.mean((tri_area - tri_area_max) ** 2)                       # (N, V, ntri_masked)
        #
        # # Compute abs residual
        # abs_residuals = torch.mean(torch.abs(tri_area - tri_area_max), dim=2, keepdim=True)     # (N, V, 1)
        #
        # return loss, abs_residuals

        loss = 0
        abs_residuals_mean = [[] for i in range(V)]
        for i in range(V - 1):
            for j in range(i + 1, V):
                # Compute loss
                cur_loss = F.mse_loss(tri_area[:, i, :], tri_area[:, j, :], reduction='none')   # (N, ntri_masked)
                loss += torch.mean(cur_loss)

                # Compute abs residual
                abs_residual = torch.abs(tri_area[:, i, :] - tri_area[:, j, :])                 # (N, ntri_masked)
                abs_residual_mean = torch.mean(abs_residual, dim=1, keepdim=True)               # (N, 1)
                abs_residuals_mean[i].append(abs_residual_mean)
                abs_residuals_mean[j].append(abs_residual_mean)
        abs_residuals_mean = [torch.mean(torch.stack(abs_residuals_mean[i], dim=0), dim=0) for i in range(V)]
        abs_residuals = torch.stack(abs_residuals_mean, dim=1)                                  # (N, V, 1)

        return loss / (V * (V - 1.0) / 2.0), abs_residuals
//  [python]
def area_loss(self, vert):
        N, V, nver, _ = vert.shape
        vert = vert * 1e-3

        # Compute area per view
        tri_mask = self.bfm.face_region_tri_mask.ravel()
        pt0 = vert[:, :, self.bfm.model['tri'][tri_mask, 2], :]                 # (N, V, ntri_masked, 3)
        pt1 = vert[:, :, self.bfm.model['tri'][tri_mask, 1], :]                 # (N, V, ntri_masked, 3)
        pt2 = vert[:, :, self.bfm.model['tri'][tri_mask, 0], :]                 # (N, V, ntri_masked, 3)
        tri_area = torch.linalg.norm(torch.cross(pt0 - pt1, pt0 - pt2, dim=3), dim=3)  # (N, V, ntri_masked)
        # tri_area_max, _ = torch.max(tri_area, dim=1, keepdim=True)
        # tri_area_max = tri_area_max.detach()                                    # (N, 1, ntri_masked)
        #
        # # Compute loss
        # loss = torch.mean((tri_area - tri_area_max) ** 2)                       # (N, V, ntri_masked)
        #
        # # Compute abs residual
        # abs_residuals = torch.mean(torch.abs(tri_area - tri_area_max), dim=2, keepdim=True)     # (N, V, 1)
        #
        # return loss, abs_residuals

        loss = 0
        abs_residuals_mean = [[] for i in range(V)]
        for i in range(V - 1):
            for j in range(i + 1, V):
                # Compute loss
                cur_loss = F.mse_loss(tri_area[:, i, :], tri_area[:, j, :], reduction='none')   # (N, ntri_masked)
                loss += torch.mean(cur_loss)

                # Compute abs residual
                abs_residual = torch.abs(tri_area[:, i, :] - tri_area[:, j, :])                 # (N, ntri_masked)
                abs_residual_mean = torch.mean(abs_residual, dim=1, keepdim=True)               # (N, 1)
                abs_residuals_mean[i].append(abs_residual_mean)
                abs_residuals_mean[j].append(abs_residual_mean)
        abs_residuals_mean = [torch.mean(torch.stack(abs_residuals_mean[i], dim=0), dim=0) for i in range(V)]
        abs_residuals = torch.stack(abs_residuals_mean, dim=1)                                  # (N, V, 1)

        return loss / (V * (V - 1.0) / 2.0), abs_residuals
//  [python]
def test_float_quant_compare_per_tensor(self):
        for qengine in supported_qengines:
            with override_quantized_engine(qengine):
                torch.manual_seed(42)
                my_model = ModelMultipleOps().to(torch.float32)
                my_model.eval()
                calib_data = torch.rand(1024, 3, 15, 15, dtype=torch.float32)
                eval_data = torch.rand(1, 3, 15, 15, dtype=torch.float32)
                out_ref = my_model(eval_data)
                qModel = torch.quantization.QuantWrapper(my_model)
                qModel.eval()
                qModel.qconfig = torch.quantization.default_qconfig
                torch.quantization.fuse_modules(qModel.module, [['conv1', 'bn1', 'relu1']], inplace=True)
                torch.quantization.prepare(qModel, inplace=True)
                qModel(calib_data)
                torch.quantization.convert(qModel, inplace=True)
                out_q = qModel(eval_data)
                SQNRdB = 20 * torch.log10(torch.norm(out_ref) / torch.norm(out_ref - out_q))
                # Quantized model output should be close to floating point model output numerically
                # Setting target SQNR to be 30 dB so that relative error is 1e-3 below the desired
                # output
                self.assertGreater(SQNRdB, 30, msg='Quantized model numerics diverge from float, expect SQNR > 30 dB')
//  [python]
def test_float_quant_compare_per_tensor(self):
        for qengine in supported_qengines:
            with override_quantized_engine(qengine):
                torch.manual_seed(42)
                my_model = ModelMultipleOps().to(torch.float32)
                my_model.eval()
                calib_data = torch.rand(1024, 3, 15, 15, dtype=torch.float32)
                eval_data = torch.rand(1, 3, 15, 15, dtype=torch.float32)
                out_ref = my_model(eval_data)
                qModel = torch.quantization.QuantWrapper(my_model)
                qModel.eval()
                qModel.qconfig = torch.quantization.default_qconfig
                torch.quantization.fuse_modules(qModel.module, [['conv1', 'bn1', 'relu1']], inplace=True)
                torch.quantization.prepare(qModel, inplace=True)
                qModel(calib_data)
                torch.quantization.convert(qModel, inplace=True)
                out_q = qModel(eval_data)
                SQNRdB = 20 * torch.log10(torch.linalg.norm(out_ref) / torch.linalg.norm(out_ref - out_q))
                # Quantized model output should be close to floating point model output numerically
                # Setting target SQNR to be 30 dB so that relative error is 1e-3 below the desired
                # output
                self.assertGreater(SQNRdB, 30, msg='Quantized model numerics diverge from float, expect SQNR > 30 dB')
//  [python]
def truncated_eig_symlobpcg(M, chi, abs_tol=1.0e-14, rel_tol=None, keep_multiplets=False, \
    eps_multiplet=1.0e-8, verbosity=0):
    r"""
    :param M: symmetric matrix of dimensions :math:`N \times N`
    :param chi: desired maximal rank :math:`\chi`
    :param abs_tol: absolute tolerance on minimal(in magnitude) eigenvalue 
    :param rel_tol: relative tolerance on minimal(in magnitude) eigenvalue
    :param keep_multiplets: truncate spectrum down to last complete multiplet
    :param eps_multiplet: allowed splitting within multiplet
    :param verbosity: logging verbosity
    :type M: torch.tensor
    :type chi: int
    :type abs_tol: float
    :type rel_tol: float
    :type keep_multiplets: bool
    :type eps_multiplet: float
    :type verbosity: int
    :return: leading :math:`\chi` eigenvalues D and eigenvectors U
    :rtype: torch.tensor, torch.tensor

    Returns leading (by magnitude) :math:`\chi` eigenpairs of a matrix M, where M is 
    a symmetric matrix :math:`M=M^T`, by computing the partial symmetric decomposition 
    :math:`MM= UDU^T` up to rank :math:`\chi`. The decomposition is computed by LOBPCG.

    Returned tensors have dimensions 

    .. math:: dim(D)=(\chi),\ dim(U)=(N,\chi)
    """
    # (optional) verify hermicity
    M_asymm_norm= torch.norm(M-M.t())
    assert M_asymm_norm/torch.abs(M).max() < 1.0e-8, "M is not symmetric"

    MM= M@M
    D2, U= SYMLOBPCG.apply(MM, chi+int(keep_multiplets), min(abs_tol,eps_multiplet))

    # find multiplets
    m=[]
    l=0
    for i in range(chi):
        l+=1
        g=D2[i]-D2[i+1]
        if g>eps_multiplet:
            m.append(l)
            l=0
            if D2[i+1]<abs_tol: break

    mixed_spec= U.t() @ M @ U
    D= torch.diag(mixed_spec)
    i=0
    for ml in m:
        if ml>1:
            Dml, Uml= torch.symeig(mixed_spec[i:i+ml,i:i+ml], eigenvectors=True)
            U[:,i:i+ml]= U[:,i:i+ml] @ Uml
            D[i:i+ml]= Dml
        i+= ml

    # estimate the chi_new 
    chi_new= chi
    if keep_multiplets:
        # regularize by discarding small values
        gaps=torch.abs(D2.clone().detach())
        # S[S < abs_tol]= 0.
        gaps[gaps < abs_tol]= 0.
        # compute gaps and normalize by the largest sing. value. Introduce cutoff
        # for handling vanishing values set to exact zero
        gaps=(gaps[:len(D2)-1]-torch.abs(D2[1:len(D2)]))/(gaps[:len(D2)-1]+1.0e-16)
        gaps[gaps > 1.0]= 0.

        if gaps[chi-1] < eps_multiplet:
            # the chi is within the multiplet - find the largest chi_new < chi
            # such that the complete multiplets are preserved
            for i in range(chi-1,-1,-1):
                if gaps[i] > eps_multiplet:
                    chi_new= i
                    break

        Dt = D[:chi].clone()
        Dt[chi_new+1:]=0.

        Ut = U[:, :Dt.shape[0]].clone()
        Ut[:, chi_new+1:]=0.
        return Dt, Ut

    return D, U
//  [python]
def truncated_eig_symlobpcg(M, chi, abs_tol=1.0e-14, rel_tol=None, keep_multiplets=False, \
    eps_multiplet=1.0e-8, verbosity=0):
    r"""
    :param M: symmetric matrix of dimensions :math:`N \times N`
    :param chi: desired maximal rank :math:`\chi`
    :param abs_tol: absolute tolerance on minimal(in magnitude) eigenvalue 
    :param rel_tol: relative tolerance on minimal(in magnitude) eigenvalue
    :param keep_multiplets: truncate spectrum down to last complete multiplet
    :param eps_multiplet: allowed splitting within multiplet
    :param verbosity: logging verbosity
    :type M: torch.tensor
    :type chi: int
    :type abs_tol: float
    :type rel_tol: float
    :type keep_multiplets: bool
    :type eps_multiplet: float
    :type verbosity: int
    :return: leading :math:`\chi` eigenvalues D and eigenvectors U
    :rtype: torch.tensor, torch.tensor

    Returns leading (by magnitude) :math:`\chi` eigenpairs of a matrix M, where M is 
    a symmetric matrix :math:`M=M^T`, by computing the partial symmetric decomposition 
    :math:`MM= UDU^T` up to rank :math:`\chi`. The decomposition is computed by LOBPCG.

    Returned tensors have dimensions 

    .. math:: dim(D)=(\chi),\ dim(U)=(N,\chi)
    """
    # (optional) verify hermicity
    M_asymm_norm= torch.linalg.norm(M-M.t())
    assert M_asymm_norm/torch.abs(M).max() < 1.0e-8, "M is not symmetric"

    MM= M@M
    D2, U= SYMLOBPCG.apply(MM, chi+int(keep_multiplets), min(abs_tol,eps_multiplet))

    # find multiplets
    m=[]
    l=0
    for i in range(chi):
        l+=1
        g=D2[i]-D2[i+1]
        if g>eps_multiplet:
            m.append(l)
            l=0
            if D2[i+1]<abs_tol: break

    mixed_spec= U.t() @ M @ U
    D= torch.diag(mixed_spec)
    i=0
    for ml in m:
        if ml>1:
            Dml, Uml= torch.symeig(mixed_spec[i:i+ml,i:i+ml], eigenvectors=True)
            U[:,i:i+ml]= U[:,i:i+ml] @ Uml
            D[i:i+ml]= Dml
        i+= ml

    # estimate the chi_new 
    chi_new= chi
    if keep_multiplets:
        # regularize by discarding small values
        gaps=torch.abs(D2.clone().detach())
        # S[S < abs_tol]= 0.
        gaps[gaps < abs_tol]= 0.
        # compute gaps and normalize by the largest sing. value. Introduce cutoff
        # for handling vanishing values set to exact zero
        gaps=(gaps[:len(D2)-1]-torch.abs(D2[1:len(D2)]))/(gaps[:len(D2)-1]+1.0e-16)
        gaps[gaps > 1.0]= 0.

        if gaps[chi-1] < eps_multiplet:
            # the chi is within the multiplet - find the largest chi_new < chi
            # such that the complete multiplets are preserved
            for i in range(chi-1,-1,-1):
                if gaps[i] > eps_multiplet:
                    chi_new= i
                    break

        Dt = D[:chi].clone()
        Dt[chi_new+1:]=0.

        Ut = U[:, :Dt.shape[0]].clone()
        Ut[:, chi_new+1:]=0.
        return Dt, Ut

    return D, U
//  [python]
def test_float_quant_compare_per_channel(self):
        # Test for per-channel Quant
        torch.manual_seed(67)
        my_model = ModelMultipleOps().to(torch.float32)
        my_model.eval()
        calib_data = torch.rand(2048, 3, 15, 15, dtype=torch.float32)
        eval_data = torch.rand(10, 3, 15, 15, dtype=torch.float32)
        out_ref = my_model(eval_data)
        q_model = torch.quantization.QuantWrapper(my_model)
        q_model.eval()
        q_model.qconfig = torch.quantization.default_per_channel_qconfig
        torch.quantization.fuse_modules(q_model.module, [['conv1', 'bn1', 'relu1']], inplace=True)
        torch.quantization.prepare(q_model)
        q_model(calib_data)
        torch.quantization.convert(q_model)
        out_q = q_model(eval_data)
        SQNRdB = 20 * torch.log10(torch.norm(out_ref) / torch.norm(out_ref - out_q))
        # Quantized model output should be close to floating point model output numerically
        # Setting target SQNR to be 35 dB
        self.assertGreater(SQNRdB, 35, msg='Quantized model numerics diverge from float, expect SQNR > 35 dB')
//  [python]
def test_float_quant_compare_per_channel(self):
        # Test for per-channel Quant
        torch.manual_seed(67)
        my_model = ModelMultipleOps().to(torch.float32)
        my_model.eval()
        calib_data = torch.rand(2048, 3, 15, 15, dtype=torch.float32)
        eval_data = torch.rand(10, 3, 15, 15, dtype=torch.float32)
        out_ref = my_model(eval_data)
        q_model = torch.quantization.QuantWrapper(my_model)
        q_model.eval()
        q_model.qconfig = torch.quantization.default_per_channel_qconfig
        torch.quantization.fuse_modules(q_model.module, [['conv1', 'bn1', 'relu1']], inplace=True)
        torch.quantization.prepare(q_model)
        q_model(calib_data)
        torch.quantization.convert(q_model)
        out_q = q_model(eval_data)
        SQNRdB = 20 * torch.log10(torch.linalg.norm(out_ref) / torch.linalg.norm(out_ref - out_q))
        # Quantized model output should be close to floating point model output numerically
        # Setting target SQNR to be 35 dB
        self.assertGreater(SQNRdB, 35, msg='Quantized model numerics diverge from float, expect SQNR > 35 dB')
//  [python]
training_size = 1000
r = torch.linspace(0.5, 2 * sigma, steps=training_size, requires_grad=True)

# Create a bunch of vectors that point along positive-x
drs = torch.outer(r, torch.tensor([1.0, 0, 0]))
norms = torch.norm(drs, dim=1).reshape(-1, 1)
//  [python]
training_size = 1000
r = torch.linspace(0.5, 2 * sigma, steps=training_size, requires_grad=True)

# Create a bunch of vectors that point along positive-x
drs = torch.outer(r, torch.tensor([1.0, 0, 0]))
norms = torch.linalg.norm(drs, dim=1).reshape(-1, 1)
//  [python]
def make_prediction(model, drs):
    norms = torch.norm(drs, dim=1).reshape(-1, 1)
    energies = model(norms)

    network_derivs = vmap(jacrev(model))(norms).squeeze(-1)
    forces = -network_derivs * drs / norms
    return energies, forces
//  [python]
def make_prediction(model, drs):
    norms = torch.linalg.norm(drs, dim=1).reshape(-1, 1)
    energies = model(norms)

    network_derivs = vmap(jacrev(model))(norms).squeeze(-1)
    forces = -network_derivs * drs / norms
    return energies, forces
//  [python]
def max_eigenvalue(mat: torch.Tensor, eps=1e-6, dup=5):
    n = mat.size(0)
    tot = 0.
    tim = 0
    with torch.no_grad():
        for i in range(dup):
            now = None
            v = torch.randn(n, dtype=torch.double)
            while True:
                tim += 1
                Av = torch.matmul(mat, v)
                vlen = torch.norm(Av)
                tmp = vlen / torch.norm(v)
                if now is None or abs(tmp - now) / now > eps:
                    now = tmp
                else:
                    now = tmp
                    break
                v = Av * (1.0 / vlen)
            tot += now.item()
    return tot / dup
//  [python]
def max_eigenvalue(mat: torch.Tensor, eps=1e-6, dup=5):
    n = mat.size(0)
    tot = 0.
    tim = 0
    with torch.no_grad():
        for i in range(dup):
            now = None
            v = torch.randn(n, dtype=torch.double)
            while True:
                tim += 1
                Av = torch.matmul(mat, v)
                vlen = torch.linalg.norm(Av)
                tmp = vlen / torch.linalg.norm(v)
                if now is None or abs(tmp - now) / now > eps:
                    now = tmp
                else:
                    now = tmp
                    break
                v = Av * (1.0 / vlen)
            tot += now.item()
    return tot / dup
//  [python]
def max_eigenvalue(self, eps=1e-6, dup=5):
        n = self.n
        tot = 0.
        tim = 0
        with torch.no_grad():
            for i in range(dup):
                now = None
                v = torch.randn(n, dtype=torch.double)
                while True:
                    tim += 1
                    Av = self.mul_vec(v)
                    vlen = torch.norm(Av)
                    tmp = vlen / torch.norm(v)
                    if now is None or abs(tmp - now) / now > eps:
                        now = tmp
                    else:
                        now = tmp
                        break
                    v = Av * (1.0 / vlen)
                tot += now.item()
        self.max_eigen = tot / dup
        return self.max_eigen
//  [python]
def max_eigenvalue(self, eps=1e-6, dup=5):
        n = self.n
        tot = 0.
        tim = 0
        with torch.no_grad():
            for i in range(dup):
                now = None
                v = torch.randn(n, dtype=torch.double)
                while True:
                    tim += 1
                    Av = self.mul_vec(v)
                    vlen = torch.linalg.norm(Av)
                    tmp = vlen / torch.linalg.norm(v)
                    if now is None or abs(tmp - now) / now > eps:
                        now = tmp
                    else:
                        now = tmp
                        break
                    v = Av * (1.0 / vlen)
                tot += now.item()
        self.max_eigen = tot / dup
        return self.max_eigen