jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def vect2euler(self, xyz):
        yz_norm = torch.norm(xyz[...,1:3], dim=-1)
        e_x = torch.atan2(-xyz[...,1], xyz[...,2])
        e_y = torch.atan2(xyz[...,0], yz_norm)
        e_z = torch.zeros_like(e_y)
        e_xyz = torch.stack([e_x, e_y, e_z], dim=-1)
        return e_xyz
//  [python]
def vect2euler(self, xyz):
        yz_norm = torch.linalg.norm(xyz[...,1:3], dim=-1)
        e_x = torch.atan2(-xyz[...,1], xyz[...,2])
        e_y = torch.atan2(xyz[...,0], yz_norm)
        e_z = torch.zeros_like(e_y)
        e_xyz = torch.stack([e_x, e_y, e_z], dim=-1)
        return e_xyz
//  [python]
def test_fake_quant_true_quant_compare(self):
        for qengine in supported_qengines:
            with override_quantized_engine(qengine):
                torch.manual_seed(67)
                my_model = ModelMultipleOpsNoAvgPool().to(torch.float32)
                calib_data = torch.rand(2048, 3, 15, 15, dtype=torch.float32)
                eval_data = torch.rand(10, 3, 15, 15, dtype=torch.float32)
                my_model.eval()
                out_ref = my_model(eval_data)
                fq_model = torch.quantization.QuantWrapper(my_model)
                fq_model.train()
                fq_model.qconfig = torch.quantization.default_qat_qconfig
                torch.quantization.fuse_modules(fq_model.module, [['conv1', 'bn1', 'relu1']], inplace=True)
                torch.quantization.prepare_qat(fq_model)
                fq_model.eval()
                fq_model.apply(torch.quantization.disable_fake_quant)
                fq_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
                fq_model(calib_data)
                fq_model.apply(torch.quantization.enable_fake_quant)
                fq_model.apply(torch.quantization.disable_observer)
                out_fq = fq_model(eval_data)
                SQNRdB = 20 * torch.log10(torch.norm(out_ref) / torch.norm(out_ref - out_fq))
                # 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')
                torch.quantization.convert(fq_model)
                out_q = fq_model(eval_data)
                SQNRdB = 20 * torch.log10(torch.norm(out_fq) / (torch.norm(out_fq - out_q) + 1e-10))
                self.assertGreater(SQNRdB, 60, msg='Fake quant and true quant numerics diverge, expect SQNR > 60 dB')
//  [python]
def test_fake_quant_true_quant_compare(self):
        for qengine in supported_qengines:
            with override_quantized_engine(qengine):
                torch.manual_seed(67)
                my_model = ModelMultipleOpsNoAvgPool().to(torch.float32)
                calib_data = torch.rand(2048, 3, 15, 15, dtype=torch.float32)
                eval_data = torch.rand(10, 3, 15, 15, dtype=torch.float32)
                my_model.eval()
                out_ref = my_model(eval_data)
                fq_model = torch.quantization.QuantWrapper(my_model)
                fq_model.train()
                fq_model.qconfig = torch.quantization.default_qat_qconfig
                torch.quantization.fuse_modules(fq_model.module, [['conv1', 'bn1', 'relu1']], inplace=True)
                torch.quantization.prepare_qat(fq_model)
                fq_model.eval()
                fq_model.apply(torch.quantization.disable_fake_quant)
                fq_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
                fq_model(calib_data)
                fq_model.apply(torch.quantization.enable_fake_quant)
                fq_model.apply(torch.quantization.disable_observer)
                out_fq = fq_model(eval_data)
                SQNRdB = 20 * torch.log10(torch.linalg.norm(out_ref) / torch.linalg.norm(out_ref - out_fq))
                # 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')
                torch.quantization.convert(fq_model)
                out_q = fq_model(eval_data)
                SQNRdB = 20 * torch.log10(torch.linalg.norm(out_fq) / (torch.linalg.norm(out_fq - out_q) + 1e-10))
                self.assertGreater(SQNRdB, 60, msg='Fake quant and true quant numerics diverge, expect SQNR > 60 dB')
//  [python]
def apply_kmeans(self, spk_vector, feature_last=False, iter_clustering=100):
        """
        Args:
            spk_vector: (batch_size, n_sources, latent_dim, T) or (batch_size, T, n_sources, latent_dim)
        Returns:
            spk_vector: (batch_size, n_clusters, latent_dim, T), where n_clusters = n_sources
        """
        if not feature_last:
            spk_vector = spk_vector.permute(0, 3, 1, 2).contiguous() # (batch_size, T, n_sources, latent_dim)

        n_sources = self.n_sources

        for idx in range(iter_clustering):
            centroids = spk_vector.mean(dim=1, keepdim=True) # (batch_size, 1, n_clusters, latent_dim)
            distance = torch.norm(spk_vector.unsqueeze(dim=3) - centroids.unsqueeze(dim=2), dim=4) # (batch_size, T, n_sources, n_clusters)
            cluster_idx = torch.argmin(distance, dim=3) # (batch_size, T, n_sources)
            mask = torch.eye(n_sources)[cluster_idx] # (batch_size, T, n_sources, n_clusters)
            mask = mask.to(spk_vector.device)
            spk_vector = torch.sum(mask.unsqueeze(dim=4) * spk_vector.unsqueeze(dim=3), dim=2) # (batch_size, T, n_clusters, latent_dim)

        if not feature_last:    
            spk_vector = spk_vector.permute(0, 2, 3, 1).contiguous() # (batch_size, n_clusters, latent_dim, T)
        
        return spk_vector
//  [python]
def apply_kmeans(self, spk_vector, feature_last=False, iter_clustering=100):
        """
        Args:
            spk_vector: (batch_size, n_sources, latent_dim, T) or (batch_size, T, n_sources, latent_dim)
        Returns:
            spk_vector: (batch_size, n_clusters, latent_dim, T), where n_clusters = n_sources
        """
        if not feature_last:
            spk_vector = spk_vector.permute(0, 3, 1, 2).contiguous() # (batch_size, T, n_sources, latent_dim)

        n_sources = self.n_sources

        for idx in range(iter_clustering):
            centroids = spk_vector.mean(dim=1, keepdim=True) # (batch_size, 1, n_clusters, latent_dim)
            distance = torch.linalg.norm(spk_vector.unsqueeze(dim=3) - centroids.unsqueeze(dim=2), dim=4) # (batch_size, T, n_sources, n_clusters)
            cluster_idx = torch.argmin(distance, dim=3) # (batch_size, T, n_sources)
            mask = torch.eye(n_sources)[cluster_idx] # (batch_size, T, n_sources, n_clusters)
            mask = mask.to(spk_vector.device)
            spk_vector = torch.sum(mask.unsqueeze(dim=4) * spk_vector.unsqueeze(dim=3), dim=2) # (batch_size, T, n_clusters, latent_dim)

        if not feature_last:    
            spk_vector = spk_vector.permute(0, 2, 3, 1).contiguous() # (batch_size, n_clusters, latent_dim, T)
        
        return spk_vector
//  [python]
def line_search(update, x0, g0, g, on=True):
    """
    Instead of solving for the best step size to use exactly, we use a fast line search algorithm
    to find an okay step size, so that compute can be spent on computing the update itself rather
    than the step size.

    `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) #for wolfe search only

    def phi(s, store=True):
        """takes in step size alpha being tried, produces the next x_est with it,
        and returns what we want to minimize, i.e. norm of g(x_est)"""
        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, on=True):
    """
    Instead of solving for the best step size to use exactly, we use a fast line search algorithm
    to find an okay step size, so that compute can be spent on computing the update itself rather
    than the step size.

    `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.norm(x0) / torch.norm(update) #for wolfe search only

    def phi(s, store=True):
        """takes in step size alpha being tried, produces the next x_est with it,
        and returns what we want to minimize, i.e. norm of g(x_est)"""
        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 get_near_sim_part_neighbours(self, label, pos, r=0.07):
        with torch.no_grad():
            # pp_label_sim.size = bz x N x N
            pp_label_sim = torch.matmul(label, label.transpose(1, 2)) / \
                           torch.clamp(torch.norm(label, dim=2, keepdim=True) * torch.norm(label.transpose(1, 2), dim=1,
                                                                                           keepdim=True), min=1e-9)

            pp_sampling_dist = pp_label_sim / torch.clamp(torch.sum(pp_label_sim, dim=-1, keepdim=True), min=1e-9)
            bz, N = pp_sampling_dist.size(0), pp_sampling_dist.size(1)
            pp_sampling_dist_expand = pp_sampling_dist.contiguous().view(bz * N, -1).contiguous()
            sampled_neis_expand = torch.multinomial(pp_sampling_dist_expand, self.pp_sim_k, replacement=True)
            sampled_neis = sampled_neis_expand.contiguous().view(bz, N, -1).contiguous()
        return sampled_neis
//  [python]
def get_near_sim_part_neighbours(self, label, pos, r=0.07):
        with torch.no_grad():
            # pp_label_sim.size = bz x N x N
            pp_label_sim = torch.matmul(label, label.transpose(1, 2)) / \
                           torch.clamp(torch.linalg.norm(label, dim=2, keepdim=True) * torch.linalg.norm(label.transpose(1, 2), dim=1,
                                                                                           keepdim=True), min=1e-9)

            pp_sampling_dist = pp_label_sim / torch.clamp(torch.sum(pp_label_sim, dim=-1, keepdim=True), min=1e-9)
            bz, N = pp_sampling_dist.size(0), pp_sampling_dist.size(1)
            pp_sampling_dist_expand = pp_sampling_dist.contiguous().view(bz * N, -1).contiguous()
            sampled_neis_expand = torch.multinomial(pp_sampling_dist_expand, self.pp_sim_k, replacement=True)
            sampled_neis = sampled_neis_expand.contiguous().view(bz, N, -1).contiguous()
        return sampled_neis
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        
        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
        (batch_size, n_frames, n_channels) 
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        # Pass the input through the LSTM layers and retrieve all outputs, the final hidden state
        # and the final cell state.
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        
        # We take only the hidden state of the last layer
        embeds_raw = self.relu(self.linear(hidden[-1]))
        
        # L2-normalize it
        embeds = embeds_raw / (torch.norm(embeds_raw, dim=1, keepdim=True) + 1e-5)        

        return embeds
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        
        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
        (batch_size, n_frames, n_channels) 
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        # Pass the input through the LSTM layers and retrieve all outputs, the final hidden state
        # and the final cell state.
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        
        # We take only the hidden state of the last layer
        embeds_raw = self.relu(self.linear(hidden[-1]))
        
        # L2-normalize it
        embeds = embeds_raw / (torch.linalg.norm(embeds_raw, dim=1, keepdim=True) + 1e-5)        

        return embeds
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.linalg.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def compute_consistent_plane_frame(normal):
    # Input:  normal is Bx3
    # Returns: x_axis, y_axis, both of dimension Bx3
    device = normal.get_device()
    batch_size, _ = normal.size()
    candidate_axes = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] # Actually, 2 should be enough. This may still cause singularity TODO!!!
    y_axes = []
    for tmp_axis in candidate_axes:
        torch_axis = torch.FloatTensor(tmp_axis).to(device).unsqueeze(0)
        y_axes.append(torch.cross(normal, torch_axis.expand(batch_size, 3)))
    y_axes = torch.stack(y_axes, dim=0) # QxBx3
    y_axes_norm = torch.norm(y_axes, dim=2) # QxB
    # choose the axis with largest norm
    y_axes_chosen_idx = torch.argmax(y_axes_norm, dim=0) # B
    y_axes_chosen_idx = y_axes_chosen_idx.view(1, batch_size, 1).expand(1, batch_size, 3)
    # y_axes_chosen[b, :] = y_axes[y_axes_chosen_idx[b], b, :]
    y_axes = torch.gather(y_axes, 0, y_axes_chosen_idx).squeeze(0)
    y_axes = torch.nn.functional.normalize(y_axes, p=2, dim=1, eps=1e-12)
    x_axes = torch.cross(y_axes, normal) # Bx3
    return x_axes, y_axes
//  [python]
def compute_consistent_plane_frame(normal):
    # Input:  normal is Bx3
    # Returns: x_axis, y_axis, both of dimension Bx3
    device = normal.get_device()
    batch_size, _ = normal.size()
    candidate_axes = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] # Actually, 2 should be enough. This may still cause singularity TODO!!!
    y_axes = []
    for tmp_axis in candidate_axes:
        torch_axis = torch.FloatTensor(tmp_axis).to(device).unsqueeze(0)
        y_axes.append(torch.cross(normal, torch_axis.expand(batch_size, 3)))
    y_axes = torch.stack(y_axes, dim=0) # QxBx3
    y_axes_norm = torch.linalg.norm(y_axes, dim=2) # QxB
    # choose the axis with largest norm
    y_axes_chosen_idx = torch.argmax(y_axes_norm, dim=0) # B
    y_axes_chosen_idx = y_axes_chosen_idx.view(1, batch_size, 1).expand(1, batch_size, 3)
    # y_axes_chosen[b, :] = y_axes[y_axes_chosen_idx[b], b, :]
    y_axes = torch.gather(y_axes, 0, y_axes_chosen_idx).squeeze(0)
    y_axes = torch.nn.functional.normalize(y_axes, p=2, dim=1, eps=1e-12)
    x_axes = torch.cross(y_axes, normal) # Bx3
    return x_axes, y_axes
//  [python]
def get_dihedral(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from:
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Inputs:
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
        * c4: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

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

    return torch.atan2( ( (torch.linalg.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )
//  [python]
def similarity_matrix(self, embeds):
        """
        Computes the similarity matrix according the section 2.1 of GE2E.

        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
        utterances_per_speaker, embedding_size)
        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
        utterances_per_speaker, speakers_per_batch)
        """
        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
        
        # Inclusive centroids (1 per speaker). Cloning is needed for reverse differentiation
        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
        centroids_incl = centroids_incl.clone() / (torch.norm(centroids_incl, dim=2, keepdim=True) + 1e-5)

        # Exclusive centroids (1 per utterance)
        centroids_excl = (torch.sum(embeds, dim=1, keepdim=True) - embeds)
        centroids_excl /= (utterances_per_speaker - 1)
        centroids_excl = centroids_excl.clone() / (torch.norm(centroids_excl, dim=2, keepdim=True) + 1e-5)

        # Similarity matrix. The cosine similarity of already 2-normed vectors is simply the dot
        # product of these vectors (which is just an element-wise multiplication reduced by a sum).
        # We vectorize the computation for efficiency.
        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker,
                                 speakers_per_batch).to(self.loss_device)
        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=np.int)
        for j in range(speakers_per_batch):
            mask = np.where(mask_matrix[j])[0]
            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
        
        ## Even more vectorized version (slower maybe because of transpose)
        # sim_matrix2 = torch.zeros(speakers_per_batch, speakers_per_batch, utterances_per_speaker
        #                           ).to(self.loss_device)
        # eye = np.eye(speakers_per_batch, dtype=np.int)
        # mask = np.where(1 - eye)
        # sim_matrix2[mask] = (embeds[mask[0]] * centroids_incl[mask[1]]).sum(dim=2)
        # mask = np.where(eye)
        # sim_matrix2[mask] = (embeds * centroids_excl).sum(dim=2)
        # sim_matrix2 = sim_matrix2.transpose(1, 2)
        
        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
        return sim_matrix
//  [python]
def similarity_matrix(self, embeds):
        """
        Computes the similarity matrix according the section 2.1 of GE2E.

        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
        utterances_per_speaker, embedding_size)
        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
        utterances_per_speaker, speakers_per_batch)
        """
        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
        
        # Inclusive centroids (1 per speaker). Cloning is needed for reverse differentiation
        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
        centroids_incl = centroids_incl.clone() / (torch.linalg.norm(centroids_incl, dim=2, keepdim=True) + 1e-5)

        # Exclusive centroids (1 per utterance)
        centroids_excl = (torch.sum(embeds, dim=1, keepdim=True) - embeds)
        centroids_excl /= (utterances_per_speaker - 1)
        centroids_excl = centroids_excl.clone() / (torch.linalg.norm(centroids_excl, dim=2, keepdim=True) + 1e-5)

        # Similarity matrix. The cosine similarity of already 2-normed vectors is simply the dot
        # product of these vectors (which is just an element-wise multiplication reduced by a sum).
        # We vectorize the computation for efficiency.
        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker,
                                 speakers_per_batch).to(self.loss_device)
        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=np.int)
        for j in range(speakers_per_batch):
            mask = np.where(mask_matrix[j])[0]
            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
        
        ## Even more vectorized version (slower maybe because of transpose)
        # sim_matrix2 = torch.zeros(speakers_per_batch, speakers_per_batch, utterances_per_speaker
        #                           ).to(self.loss_device)
        # eye = np.eye(speakers_per_batch, dtype=np.int)
        # mask = np.where(1 - eye)
        # sim_matrix2[mask] = (embeds[mask[0]] * centroids_incl[mask[1]]).sum(dim=2)
        # mask = np.where(eye)
        # sim_matrix2[mask] = (embeds * centroids_excl).sum(dim=2)
        # sim_matrix2 = sim_matrix2.transpose(1, 2)
        
        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
        return sim_matrix
//  [python]
def __call__(self, logits, onehot_labels, feature_maps):
        onehot_labels = torch.zeros(logits.size()).scatter_(
            1, onehot_labels.unsqueeze(1).cpu(), 1).to(logits.device)
        prob = F.softmax(logits, dim=1)
        correct_prob = prob * onehot_labels

        correct_prob = torch.sum(correct_prob, dim=1, keepdim=True)
        other_prob = prob * (1.0 - onehot_labels)

        if self.top_k > 1:
            topk_prob, _ = other_prob.topk(self.top_k, dim=1)
        else:
            topk_prob, _ = other_prob.max(dim=1, keepdim=True)

        diff_prob = correct_prob - topk_prob

        loss = torch.empty(0, device=logits.device)
        for feature_map in feature_maps:
            diff_grad = torch.stack([_get_grad(diff_prob[:, i], feature_map) for i in range(self.top_k)],
                                    dim=1)
            diff_gradnorm = torch.norm(diff_grad, p=self.dual_norm, dim=2)

            if self.use_approximation:
                diff_gradnorm.detach_()

            dist_to_boundary = diff_prob / (diff_gradnorm + self.eps)

            if self.loss_type == "worst_top_k":
                dist_to_boundary, _ = dist_to_boundary.min(dim=1)
            elif self.loss_type == "avg_top_k":
                dist_to_boundary = dist_to_boundary.mean(dim=1)

            loss_layer = _max_with_relu(dist_to_boundary, self.dist_lower)
            loss_layer = _max_with_relu(
                0, self.dist_upper - loss_layer) - self.dist_upper
            loss = torch.cat([loss, loss_layer])
        if self.reduce == 'mean':
            return loss.mean()
        else:
            if self.loss_type in ['worst_top_k', 'avg_top_k']:
                return loss
            return loss.mean(dim=1)
//  [python]
def __call__(self, logits, onehot_labels, feature_maps):
        onehot_labels = torch.zeros(logits.size()).scatter_(
            1, onehot_labels.unsqueeze(1).cpu(), 1).to(logits.device)
        prob = F.softmax(logits, dim=1)
        correct_prob = prob * onehot_labels

        correct_prob = torch.sum(correct_prob, dim=1, keepdim=True)
        other_prob = prob * (1.0 - onehot_labels)

        if self.top_k > 1:
            topk_prob, _ = other_prob.topk(self.top_k, dim=1)
        else:
            topk_prob, _ = other_prob.max(dim=1, keepdim=True)

        diff_prob = correct_prob - topk_prob

        loss = torch.empty(0, device=logits.device)
        for feature_map in feature_maps:
            diff_grad = torch.stack([_get_grad(diff_prob[:, i], feature_map) for i in range(self.top_k)],
                                    dim=1)
            diff_gradnorm = torch.linalg.norm(diff_grad, ord=self.dual_norm, dim=2)

            if self.use_approximation:
                diff_gradnorm.detach_()

            dist_to_boundary = diff_prob / (diff_gradnorm + self.eps)

            if self.loss_type == "worst_top_k":
                dist_to_boundary, _ = dist_to_boundary.min(dim=1)
            elif self.loss_type == "avg_top_k":
                dist_to_boundary = dist_to_boundary.mean(dim=1)

            loss_layer = _max_with_relu(dist_to_boundary, self.dist_lower)
            loss_layer = _max_with_relu(
                0, self.dist_upper - loss_layer) - self.dist_upper
            loss = torch.cat([loss, loss_layer])
        if self.reduce == 'mean':
            return loss.mean()
        else:
            if self.loss_type in ['worst_top_k', 'avg_top_k']:
                return loss
            return loss.mean(dim=1)
//  [python]
def get_angle(c1, c2, c3):
    """ Returns the angle in radians. Uses atan2 formula
        Inputs:
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2

    # dont use acos since norms involved.
    # better use atan2 formula: atan2(cross, dot) from here:
    # https://johnblackburne.blogspot.com/2012/05/angle-between-two-3d-vectors.html

    # add a minus since we want the angle in reversed order - sidechainnet issues
    return torch.atan2( torch.norm(torch.cross(u1,u2, dim=-1), dim=-1),
                        -(u1*u2).sum(dim=-1) )
//  [python]
def get_angle(c1, c2, c3):
    """ Returns the angle in radians. Uses atan2 formula
        Inputs:
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2

    # dont use acos since norms involved.
    # better use atan2 formula: atan2(cross, dot) from here:
    # https://johnblackburne.blogspot.com/2012/05/angle-between-two-3d-vectors.html

    # add a minus since we want the angle in reversed order - sidechainnet issues
    return torch.atan2( torch.linalg.norm(torch.cross(u1,u2, dim=-1), dim=-1),
                        -(u1*u2).sum(dim=-1) )
//  [python]
def _get_info(self):
        err_x = torch.norm((self.x[..., :3]-self.traj_x[...,self.step_cnt])[:,:2],dim=1)
        err_v = torch.norm((self.v[..., :3]-self.traj_v[...,self.step_cnt])[:,:2],dim=1)
        info = {
            'pos': self.x,
            'vel': self.v,
            'acc': self.acc, 
            'action': self.action, 
            'action_force': self.action_force, 
            'mass': self.mass, 
            'disturb': self.disturb,
            'decay': self.decay,
            'err_x': err_x, 
            'err_v': err_v, 
            'e': self._get_e(),
            'obs_his': torch.stack(self.obs_his, dim=0),
            'u_his': torch.stack(self.u_his, dim=0),
            'd_u_his': torch.stack(self.d_u_his, dim=0),
            # 'u_force_his': torch.ones_like(torch.stack(self.u_force_his, dim=0), device=self.device)*1000,
            'u_force_his': torch.stack(self.u_force_his, dim=0), 
            'd_u_force_his': torch.stack(self.d_u_force_his, dim=0),
            'v_his': torch.stack(self.v_his, dim=0),
            'acc_his': torch.stack(self.acc_his, dim=0),
            'd_acc_his': torch.stack(self.d_acc_his, dim=0),
            'delay': self.delay,
        }

        info['adapt_obs'] = torch.concat(
            [info['u_force_his'], info['d_u_force_his'], (info['v_his']-self.v_mean)/self.v_std, 
            (info['acc_his']-self.acc_mean)/self.acc_std, (info['d_acc_his']-self.d_acc_mean)/self.d_acc_std], 
            dim=-1
        ) # Tensor(adapt_horizon, env_num, param)]
        info['adapt_obs'] = info['adapt_obs'].permute(1,0,2)
        info['adapt_obs'] = info['adapt_obs'].reshape(self.env_num, -1)
        # info['adapt_obs'] = torch.concat([(info['acc_his'][-1]-self.acc_mean)/self.acc_std, info['u_force_his'][-1], self.gravity/9.8], dim=-1) # Tensor(env_num, 3[dim], 3[param])

        if self.delay_max > 0:
            info['action_history'] = torch.stack(self.action_history, dim=1)
        return info
//  [python]
def _get_info(self):
        err_x = torch.linalg.norm((self.x[..., :3]-self.traj_x[...,self.step_cnt])[:,:2],dim=1)
        err_v = torch.linalg.norm((self.v[..., :3]-self.traj_v[...,self.step_cnt])[:,:2],dim=1)
        info = {
            'pos': self.x,
            'vel': self.v,
            'acc': self.acc, 
            'action': self.action, 
            'action_force': self.action_force, 
            'mass': self.mass, 
            'disturb': self.disturb,
            'decay': self.decay,
            'err_x': err_x, 
            'err_v': err_v, 
            'e': self._get_e(),
            'obs_his': torch.stack(self.obs_his, dim=0),
            'u_his': torch.stack(self.u_his, dim=0),
            'd_u_his': torch.stack(self.d_u_his, dim=0),
            # 'u_force_his': torch.ones_like(torch.stack(self.u_force_his, dim=0), device=self.device)*1000,
            'u_force_his': torch.stack(self.u_force_his, dim=0), 
            'd_u_force_his': torch.stack(self.d_u_force_his, dim=0),
            'v_his': torch.stack(self.v_his, dim=0),
            'acc_his': torch.stack(self.acc_his, dim=0),
            'd_acc_his': torch.stack(self.d_acc_his, dim=0),
            'delay': self.delay,
        }

        info['adapt_obs'] = torch.concat(
            [info['u_force_his'], info['d_u_force_his'], (info['v_his']-self.v_mean)/self.v_std, 
            (info['acc_his']-self.acc_mean)/self.acc_std, (info['d_acc_his']-self.d_acc_mean)/self.d_acc_std], 
            dim=-1
        ) # Tensor(adapt_horizon, env_num, param)]
        info['adapt_obs'] = info['adapt_obs'].permute(1,0,2)
        info['adapt_obs'] = info['adapt_obs'].reshape(self.env_num, -1)
        # info['adapt_obs'] = torch.concat([(info['acc_his'][-1]-self.acc_mean)/self.acc_std, info['u_force_his'][-1], self.gravity/9.8], dim=-1) # Tensor(env_num, 3[dim], 3[param])

        if self.delay_max > 0:
            info['action_history'] = torch.stack(self.action_history, dim=1)
        return info
//  [python]
def cal_approx_hessian_vec_prod4(truncted_s, extended_Y_k_list, i, m, v_vec):
    
#     curr_S_k = torch.zeros([v_vec.shape[0],m], dtype = torch.double)#S_k_list[:,i-m:i]
#         
#     curr_Y_k = torch.zeros([v_vec.shape[0],m], dtype = torch.double)
#     
#     for k in range(m):
#         curr_S_k[:,k] = get_all_vectorized_parameters(para_list_all_epochs[i-(k+1)]) - get_all_vectorized_parameters(para_list_all_epochs[i])
#     
#         curr_Y_k[:,k] = get_all_vectorized_parameters(gradient_list_all_epoch[i-(k+1)]) - get_all_vectorized_parameters(gradient_list_all_epoch[i])
    
    
    res = torch.mm(torch.inverse(torch.mm(torch.t(truncted_s), truncted_s)), torch.mm(torch.t(truncted_s), v_vec))
    
    print(res)
    
    print(torch.norm(torch.mm(truncted_s, res) - v_vec))
    
    return torch.mm(extended_Y_k_list, res.view(-1,1)), torch.mm(truncted_s, res)
//  [python]
def cal_approx_hessian_vec_prod4(truncted_s, extended_Y_k_list, i, m, v_vec):
    
    res = torch.mm(torch.inverse(torch.mm(torch.t(truncted_s), truncted_s)), torch.mm(torch.t(truncted_s), v_vec))
    
    print(res)
    
    print(torch.linalg.norm(torch.mm(truncted_s, res) - v_vec))
    
    return torch.mm(extended_Y_k_list, res.view(-1,1)), torch.mm(truncted_s, res)
//  [python]
def init_prod(state, env, verbosity=0):
    for key,t in env.C.items():
        env.C[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.C[key][0,0]= 1.0 + 0.j if env.C[key].is_complex() else 1.0

    A= next(iter(state.sites.values()))
    # left transfer matrix
    #
    #     0      = 0     
    # i--A--j      T
    #   /\         1
    #  2  m
    #      \ 0
    #    i--A--j
    #      /
    #     2
    a = torch.einsum('meifj,maibj->eafb',(A,A.conj())).contiguous().view(\
        A.size()[1]**2, A.size()[3]**2)
    a= a/a.abs().max()
    # check symmetry
    a_asymm_norm= torch.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, 2)
    # leading eigenvector is unique 
    assert torch.abs(D[0]-D[1]) > 1.0e-8, "Leading eigenvector of T not unique"
    for key,t in env.T.items():
        env.T[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.T[key][0,0,:]= U[:,0]
//  [python]
def init_prod(state, env, verbosity=0):
    for key,t in env.C.items():
        env.C[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.C[key][0,0]= 1.0 + 0.j if env.C[key].is_complex() else 1.0

    A= next(iter(state.sites.values()))
    # left transfer matrix
    #
    #     0      = 0     
    # i--A--j      T
    #   /\         1
    #  2  m
    #      \ 0
    #    i--A--j
    #      /
    #     2
    a = torch.einsum('meifj,maibj->eafb',(A,A.conj())).contiguous().view(\
        A.size()[1]**2, A.size()[3]**2)
    a= a/a.abs().max()
    # check symmetry
    a_asymm_norm= torch.linalg.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, 2)
    # leading eigenvector is unique 
    assert torch.abs(D[0]-D[1]) > 1.0e-8, "Leading eigenvector of T not unique"
    for key,t in env.T.items():
        env.T[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.T[key][0,0,:]= U[:,0]
//  [python]
def compute_gradient_diff(curr_gradient_list, gradient_list):
    
    print("gradient diff:")
    
    for i in range(len(curr_gradient_list)):
        
        print(curr_gradient_list[i])
        
        print(gradient_list[i])
        
        
        print("gradient_dff::", i, torch.norm(curr_gradient_list[i]-gradient_list[i]))
//  [python]
def compute_gradient_diff(curr_gradient_list, gradient_list):
    
    print("gradient diff:")
    
    for i in range(len(curr_gradient_list)):
        
        print(curr_gradient_list[i])
        
        print(gradient_list[i])
        
        
        print("gradient_dff::", i, torch.linalg.norm(curr_gradient_list[i]-gradient_list[i]))
//  [python]
def RandomWalkRestart(A, rst_prob, delta=1e-4, reset=None, max_iter=50, use_torch=False, return_torch=False):
    if use_torch:
        device = torch.device("cuda:0")
    nnode = A.shape[0]
    # print nnode
    if reset is None:
        reset = np.eye(nnode)
    nsample, nnode = reset.shape
    # print nsample,nnode
    P = renorm(A)
    P = P.T
    norm_reset = renorm(reset.T)
    norm_reset = norm_reset.T
    if use_torch:
        norm_reset = torch.from_numpy(norm_reset).float().to(device)
        P = torch.from_numpy(P).float().to(device)
    Q = norm_reset

    for i in range(1, max_iter):
        # Q = gnp.garray(Q)
        # P = gnp.garray(P)
        if use_torch:
            Q_new = rst_prob * norm_reset + (1 - rst_prob) * torch.mm(Q, P)  # .as_numpy_array()
            delta = torch.norm(Q - Q_new, 2)
        else:
            Q_new = rst_prob * norm_reset + (1 - rst_prob) * np.dot(Q, P)  # .as_numpy_array()
            delta = np.linalg.norm(Q - Q_new, 'fro')
        Q = Q_new
        # print (i,Q)
        sys.stdout.flush()
        if delta < 1e-4:
            break
    if use_torch and not return_torch:
        Q = Q.cpu().numpy()
    return Q
//  [python]
def RandomWalkRestart(A, rst_prob, delta=1e-4, reset=None, max_iter=50, use_torch=False, return_torch=False):
    if use_torch:
        device = torch.device("cuda:0")
    nnode = A.shape[0]
    # print nnode
    if reset is None:
        reset = np.eye(nnode)
    nsample, nnode = reset.shape
    # print nsample,nnode
    P = renorm(A)
    P = P.T
    norm_reset = renorm(reset.T)
    norm_reset = norm_reset.T
    if use_torch:
        norm_reset = torch.from_numpy(norm_reset).float().to(device)
        P = torch.from_numpy(P).float().to(device)
    Q = norm_reset

    for i in range(1, max_iter):
        # Q = gnp.garray(Q)
        # P = gnp.garray(P)
        if use_torch:
            Q_new = rst_prob * norm_reset + (1 - rst_prob) * torch.mm(Q, P)  # .as_numpy_array()
            delta = torch.linalg.norm(Q - Q_new, 2)
        else:
            Q_new = rst_prob * norm_reset + (1 - rst_prob) * np.dot(Q, P)  # .as_numpy_array()
            delta = np.linalg.norm(Q - Q_new, 'fro')
        Q = Q_new
        # print (i,Q)
        sys.stdout.flush()
        if delta < 1e-4:
            break
    if use_torch and not return_torch:
        Q = Q.cpu().numpy()
    return Q
//  [python]
def compute_updated_parameters(old_delta, expected_delta, delta_gradient_list, old_para_list, para_list):
    
    
    new_delta = old_delta.clone()
    
    
    for i in range(len(old_para_list)):
        para_updates = para_list[i] - old_para_list[i]
        
        for j in range(20):
            
            for k in range(old_delta.shape[1]):
                new_delta[j][k] += torch.sum(delta_gradient_list[j][k][i]*para_updates)
            
    
    print("delta_diff::", torch.norm(new_delta[0:20]-expected_delta[0:20]))
//  [python]
def compute_updated_parameters(old_delta, expected_delta, delta_gradient_list, old_para_list, para_list):
    
    
    new_delta = old_delta.clone()
    
    
    for i in range(len(old_para_list)):
        para_updates = para_list[i] - old_para_list[i]
        
        for j in range(20):
            
            for k in range(old_delta.shape[1]):
                new_delta[j][k] += torch.sum(delta_gradient_list[j][k][i]*para_updates)
            
    
    print("delta_diff::", torch.linalg.norm(new_delta[0:20]-expected_delta[0:20]))
//  [python]
def init_from_ipeps_pbc(state, env, verbosity=0):
    if verbosity>0:
        print("ENV: init_from_ipeps_pbc")

    # Left-upper corner
    #
    #     i      = C--1     
    # j--A--3      0
    #   /\
    #  2  m
    #      \ i
    #    j--A--3
    #      /
    #     2
    A= next(iter(state.sites.values()))
    dimsA= A.size()
    a= torch.einsum('mijef,mijab->eafb',A,A.conj()).contiguous().view(dimsA[3]**2, dimsA[4]**2)
    a= a/a.abs().max()

    a_asymm_norm= torch.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, a.size()[0])
    a= torch.diag(D)

    env.C[env.keyC]= torch.zeros(env.chi,env.chi, dtype=env.dtype, device=env.device)
    env.C[env.keyC][:min(env.chi,dimsA[3]**2),:min(env.chi,dimsA[4]**2)]=\
       a[:min(env.chi,dimsA[3]**2),:min(env.chi,dimsA[4]**2)]

    # left transfer matrix (orientation from 1->0)
    #
    #     0      = A 0     
    # i--A--3      | T--2
    #   /\         | 1
    #  2  m
    #      \ 0
    #    i--A--3
    #      /
    #     2
    a= torch.einsum('meifg,maibc->eafbgc',(A,A.conj())).contiguous().view(dimsA[1]**2, dimsA[3]**2, dimsA[4]**2)
    a= a/a.abs().max()

    a= torch.einsum('ai,abs,bj->ijs',U,a,U.conj())
    a_asymm_norm= (a-a.permute(1,0,2).conj()).norm()
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"

    env.T[env.keyT]= torch.zeros((env.chi,env.chi,dimsA[4]**2), dtype=env.dtype, device=env.device)
    env.T[env.keyT][:min(env.chi,dimsA[1]**2),:min(env.chi,dimsA[3]**2),:]=\
        a[:min(env.chi,dimsA[1]**2),:min(env.chi,dimsA[3]**2),:]
//  [python]
def init_from_ipeps_pbc(state, env, verbosity=0):
    if verbosity>0:
        print("ENV: init_from_ipeps_pbc")

    # Left-upper corner
    #
    #     i      = C--1     
    # j--A--3      0
    #   /\
    #  2  m
    #      \ i
    #    j--A--3
    #      /
    #     2
    A= next(iter(state.sites.values()))
    dimsA= A.size()
    a= torch.einsum('mijef,mijab->eafb',A,A.conj()).contiguous().view(dimsA[3]**2, dimsA[4]**2)
    a= a/a.abs().max()

    a_asymm_norm= torch.linalg.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, a.size()[0])
    a= torch.diag(D)

    env.C[env.keyC]= torch.zeros(env.chi,env.chi, dtype=env.dtype, device=env.device)
    env.C[env.keyC][:min(env.chi,dimsA[3]**2),:min(env.chi,dimsA[4]**2)]=\
       a[:min(env.chi,dimsA[3]**2),:min(env.chi,dimsA[4]**2)]

    # left transfer matrix (orientation from 1->0)
    #
    #     0      = A 0     
    # i--A--3      | T--2
    #   /\         | 1
    #  2  m
    #      \ 0
    #    i--A--3
    #      /
    #     2
    a= torch.einsum('meifg,maibc->eafbgc',(A,A.conj())).contiguous().view(dimsA[1]**2, dimsA[3]**2, dimsA[4]**2)
    a= a/a.abs().max()

    a= torch.einsum('ai,abs,bj->ijs',U,a,U.conj())
    a_asymm_norm= torch.linalg.norm(a-a.permute(1,0,2).conj())
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"

    env.T[env.keyT]= torch.zeros((env.chi,env.chi,dimsA[4]**2), dtype=env.dtype, device=env.device)
    env.T[env.keyT][:min(env.chi,dimsA[1]**2),:min(env.chi,dimsA[3]**2),:]=\
        a[:min(env.chi,dimsA[1]**2),:min(env.chi,dimsA[3]**2),:]
//  [python]
def compute_grad_final5(para, hessian_para_prod, gradient_dual, grad_list_tensor, para_list_tensor, size1, size2, alpha, beta, is_GPU, device):
    
#     gradient_list = []
#     
#     exp_grad_list = []
    
    gradients = None
    
    exp_grad_list_full = None
    
    
    delta_grad_norm = torch.norm(hessian_para_prod)
    
    if gradient_dual is not None:
        
        hessian_para_prod += grad_list_tensor 
        
        hessian_para_prod += beta*para_list_tensor 
        
        
#         tmp = (grad_list_tensor + beta*para_list_tensor)
        
#         hessian_para_prod += tmp.view(-1,1)
        gradients = hessian_para_prod*size1
        
        gradients += (gradient_dual + beta*para_list_tensor)*size2
        
        gradients /= (size1 + size2)
        
        
#         for i in range(len(para)):
#             
#             if is_GPU:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i]).to(device)
#             else:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i])
#             
#             
#             curr_hessian_para_prod = hessian_para_prod[i]
#             
#             
#             gradients = (curr_hessian_para_prod*size1 + (gradient_dual[i] + beta*para[i])*size2)/(size1 + size2)
#         
#             
#             gradient_list.append(gradients)
        
    else:
        
        hessian_para_prod += (grad_list_tensor + beta*para_list_tensor)
        
        gradients = hessian_para_prod
        
#         for i in range(len(para)):
#             
# 
#             if is_GPU:        
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i]).to(device)
#             else:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i])
#             
#             gradients = hessian_para_prod[i]
#             
#             
#             gradient_list.append(gradients)
            
#     print(gradients.shape)

    exp_grad_list_full = gradients - beta*para
    delta_para = para - para_list_tensor
     
    
     
    tmp_res = 0
     
     
    if torch.norm(delta_para) > delta_grad_norm:
        return True, gradients, exp_grad_list_full
 
    else:
        return False, gradients, exp_grad_list_full
//  [python]
def compute_grad_final5(para, hessian_para_prod, gradient_dual, grad_list_tensor, para_list_tensor, size1, size2, alpha, beta, is_GPU, device):
    
#     gradient_list = []
#     
#     exp_grad_list = []
    
    gradients = None
    
    exp_grad_list_full = None
    
    
    delta_grad_norm = torch.linalg.norm(hessian_para_prod)
    
    if gradient_dual is not None:
        
        hessian_para_prod += grad_list_tensor 
        
        hessian_para_prod += beta*para_list_tensor 
        
        
#         tmp = (grad_list_tensor + beta*para_list_tensor)
        
#         hessian_para_prod += tmp.view(-1,1)
        gradients = hessian_para_prod*size1
        
        gradients += (gradient_dual + beta*para_list_tensor)*size2
        
        gradients /= (size1 + size2)
        
        
#         for i in range(len(para)):
#             
#             if is_GPU:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i]).to(device)
#             else:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i])
#             
#             
#             curr_hessian_para_prod = hessian_para_prod[i]
#             
#             
#             gradients = (curr_hessian_para_prod*size1 + (gradient_dual[i] + beta*para[i])*size2)/(size1 + size2)
#         
#             
#             gradient_list.append(gradients)
        
    else:
        
        hessian_para_prod += (grad_list_tensor + beta*para_list_tensor)
        
        gradients = hessian_para_prod
        
#         for i in range(len(para)):
#             
# 
#             if is_GPU:        
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i]).to(device)
#             else:
#                 hessian_para_prod[i] += (grad_list[i] + beta*para_list[i])
#             
#             gradients = hessian_para_prod[i]
#             
#             
#             gradient_list.append(gradients)
            
#     print(gradients.shape)

    exp_grad_list_full = gradients - beta*para
    delta_para = para - para_list_tensor
     
    
     
    tmp_res = 0
     
     
    if torch.linalg.norm(delta_para) > delta_grad_norm:
        return True, gradients, exp_grad_list_full
 
    else:
        return False, gradients, exp_grad_list_full
//  [python]
def forward(self, mels):
        self.lstm.flatten_parameters()
        _, (hidden, _) = self.lstm(mels.transpose(-1, -2))
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def forward(self, mels):
        self.lstm.flatten_parameters()
        _, (hidden, _) = self.lstm(mels.transpose(-1, -2))
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)