jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sym_penalty(dataset,keypoints,pred_out):
    """
    get penalty for the symmetry of human body
    :return:
    """
    loss_sym = 0
    if dataset == 'h36m':
        if keypoints.startswith('hr'):
            left_bone = [(0,4),(4,5),(5,6),(8,10),(10,11),(11,12)]
            right_bone = [(0,1),(1,2),(2,3),(8,13),(13,14),(14,15)]
        else:
            left_bone = [(0,4),(4,5),(5,6),(8,11),(11,12),(12,13)]
            right_bone = [(0,1),(1,2),(2,3),(8,14),(14,15),(15,16)]
        for (i_left,j_left),(i_right,j_right) in zip(left_bone,right_bone):
            left_part = pred_out[:,:,i_left]-pred_out[:,:,j_left]
            right_part = pred_out[:, :, i_right] - pred_out[:, :, j_right]
            loss_sym += torch.mean(torch.abs(torch.linalg.norm(left_part, dim=-1) - torch.linalg.norm(right_part, dim=-1)))
    elif dataset.startswith('STB'):
        loss_sym = 0
    return 0.01*loss_sym
//  [python]
def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
        if self.training and self.moe_jitter_eps is not None:
            hidden_states *= torch.empty_like(hidden_states).uniform_(
                1.0 - self.moe_jitter_eps, 1.0 + self.moe_jitter_eps
            )
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        weights = self.layer(hidden_states).softmax(dim=-1, dtype=torch.float32)
        top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)

        top_weights_scale = (
            torch.norm(top_weights, p=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
            if self.moe_normalize_expert_weights is not None
            else 1.0
        )
        top_weights = top_weights / top_weights_scale

        weights = weights.to(hidden_states.dtype)
        top_weights = top_weights.to(hidden_states.dtype)
        return weights, top_weights, top_experts
//  [python]
def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
        if self.training and self.moe_jitter_eps is not None:
            hidden_states *= torch.empty_like(hidden_states).uniform_(
                1.0 - self.moe_jitter_eps, 1.0 + self.moe_jitter_eps
            )
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        weights = self.layer(hidden_states).softmax(dim=-1, dtype=torch.float32)
        top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)

        top_weights_scale = (
            torch.linalg.norm(top_weights, ord=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
            if self.moe_normalize_expert_weights is not None
            else 1.0
        )
        top_weights = top_weights / top_weights_scale

        weights = weights.to(hidden_states.dtype)
        top_weights = top_weights.to(hidden_states.dtype)
        return weights, top_weights, top_experts
//  [python]
def forward(self, x, y, z):
        x = torch.norm(x)
        y = torch.norm(y, p=2, dim=(2,3), keepdim=False)
        z = torch.norm(z, p=1, dim=None, keepdim=True)
        return x, y, z
//  [python]
def forward(self, x, y, z):
        x = torch.linalg.norm(x)
        y = torch.linalg.norm(y, ord=2, dim=(2,3), keepdim=False)
        z = torch.linalg.norm(z, ord=1, keepdim=True)
        return x, y, z
//  [python]
def forward(self, x, **kwargs):
        n = torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps)
        x = x / n * self.g
        return self.fn(x, **kwargs)
//  [python]
def forward(self, x, **kwargs):
        n = torch.linalg.norm(x, dim=-1, keepdim=True).clamp(min=self.eps)
        x = x / n * self.g
        return self.fn(x, **kwargs)
//  [python]
def construct_M_logits(self, predicted_centers, object_centers):
        """ Construct the logits of tensor M given the object centers, and the predicted centers
            for a SINGLE frame.

            Note: We do NOT multiply by foreground here. This is mostly used for training

            @param predicted_centers: a [3 x H x W] torch.FloatTensor of center predictions
            @param object_centers: a [num_objects x 3] torch.FloatTensor. 
            @param fg_logits: a [3 x H x W] torch.FloatTensor of foreground logits

            @return: a [num_objects x H x W] torch.FloatTensor
        """

        # Compute expontiated distances, then apply LogSoftmax
        distances = torch.norm(predicted_centers.unsqueeze(0) - object_centers.unsqueeze(-1).unsqueeze(-1), dim=1) # Shape: [num_obj x H x W]

        return -self.config['tau'] * distances
//  [python]
def construct_M_logits(self, predicted_centers, object_centers):
        """ Construct the logits of tensor M given the object centers, and the predicted centers
            for a SINGLE frame.

            Note: We do NOT multiply by foreground here. This is mostly used for training

            @param predicted_centers: a [3 x H x W] torch.FloatTensor of center predictions
            @param object_centers: a [num_objects x 3] torch.FloatTensor. 
            @param fg_logits: a [3 x H x W] torch.FloatTensor of foreground logits

            @return: a [num_objects x H x W] torch.FloatTensor
        """

        # Compute expontiated distances, then apply LogSoftmax
        distances = torch.linalg.norm(predicted_centers.unsqueeze(0) - object_centers.unsqueeze(-1).unsqueeze(-1), dim=1) # Shape: [num_obj x H x W]

        return -self.config['tau'] * distances
//  [python]
def _attack_step(self):
        e = 1e-6
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                grad = self.optimizer.state[param]['exp_avg']
                norm_grad = torch.norm(grad)
                norm_data = torch.norm(param.detach())

                if norm_grad != 0 and not torch.isnan(norm_grad):
                    # Set lower and upper limit in change
                    limit_eps = self.adv_eps * param.detach().abs()
                    param_min = param.data - limit_eps
                    param_max = param.data + limit_eps

                    # Perturb along gradient
                    # w += (adv_lr * |w| / |grad|) * grad
                    param.data.add_(grad, alpha=(self.adv_lr * (norm_data + e) / (norm_grad + e)))

                    # Apply the limit to the change
                    param.data.clamp_(param_min, param_max)
//  [python]
def _attack_step(self):
        e = 1e-6
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                grad = self.optimizer.state[param]['exp_avg']
                norm_grad = torch.linalg.norm(grad)
                norm_data = torch.linalg.norm(param.detach())

                if norm_grad != 0 and not torch.isnan(norm_grad):
                    # Set lower and upper limit in change
                    limit_eps = self.adv_eps * param.detach().abs()
                    param_min = param.data - limit_eps
                    param_max = param.data + limit_eps

                    # Perturb along gradient
                    # w += (adv_lr * |w| / |grad|) * grad
                    param.data.add_(grad, alpha=(self.adv_lr * (norm_data + e) / (norm_grad + e)))

                    # Apply the limit to the change
                    param.data.clamp_(param_min, param_max)
//  [python]
def loss_mpjpe(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def loss_mpjpe(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def weighted_mpjpe(predicted, target, w):
    """
    Weighted mean per-joint position error (i.e. mean Euclidean distance)
    """
    assert predicted.shape == target.shape
    assert w.shape[0] == predicted.shape[0]
    return torch.mean(w * torch.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def weighted_mpjpe(predicted, target, w):
    """
    Weighted mean per-joint position error (i.e. mean Euclidean distance)
    """
    assert predicted.shape == target.shape
    assert w.shape[0] == predicted.shape[0]
    return torch.mean(w * torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def loss_2d_weighted(predicted, target, conf):
    assert predicted.shape == target.shape
    predicted_2d = predicted[:,:,:,:2]
    target_2d = target[:,:,:,:2]
    diff = (predicted_2d - target_2d) * conf
    return torch.mean(torch.norm(diff, dim=-1))
//  [python]
def loss_2d_weighted(predicted, target, conf):
    assert predicted.shape == target.shape
    predicted_2d = predicted[:,:,:,:2]
    target_2d = target[:,:,:,:2]
    diff = (predicted_2d - target_2d) * conf
    return torch.mean(torch.linalg.norm(diff, dim=-1))
//  [python]
def loss_velocity(predicted, target):
    """
    Mean per-joint velocity error (i.e. mean Euclidean distance of the 1st derivative)
    """
    assert predicted.shape == target.shape
    if predicted.shape[1]<=1:
        return torch.FloatTensor(1).fill_(0.)[0].to(predicted.device)
    velocity_predicted = predicted[:,1:] - predicted[:,:-1]
    velocity_target = target[:,1:] - target[:,:-1]
    return torch.mean(torch.norm(velocity_predicted - velocity_target, dim=-1))
//  [python]
def loss_velocity(predicted, target):
    """
    Mean per-joint velocity error (i.e. mean Euclidean distance of the 1st derivative)
    """
    assert predicted.shape == target.shape
    if predicted.shape[1]<=1:
        return torch.FloatTensor(1).fill_(0.)[0].to(predicted.device)
    velocity_predicted = predicted[:,1:] - predicted[:,:-1]
    velocity_target = target[:,1:] - target[:,:-1]
    return torch.mean(torch.linalg.norm(velocity_predicted - velocity_target, dim=-1))
//  [python]
def aux_loss(pos_y, pos_x, batch):
        edge_index = batch.edge_index
        src = edge_index[0]
        tgt = edge_index[1]
        true_bond = torch.norm(pos_y[src] - pos_y[tgt], dim=-1)
        pred_bond = torch.norm(pos_x[src] - pos_x[tgt], dim=-1)
        bond_loss = torch.mean(F.l1_loss(pred_bond, true_bond))

        nei_src_index = batch.nei_src_index.view(-1)
        nei_tgt_index = batch.nei_tgt_index
        nei_tgt_mask = batch.nei_tgt_mask

        random_tgt_index = pos_y.new_zeros(nei_tgt_index.size()).uniform_()
        random_tgt_index = torch.where(
            nei_tgt_mask, pos_y.new_zeros(nei_tgt_index.size()), random_tgt_index
        )
        random_tgt_index_sort = torch.sort(random_tgt_index, descending=True, dim=0)[1][:2]
        tgt_1, tgt_2 = random_tgt_index_sort[0].unsqueeze(0), random_tgt_index_sort[1].unsqueeze(0)

        tgt_1 = torch.gather(nei_tgt_index, 0, tgt_1).view(-1)
        tgt_2 = torch.gather(nei_tgt_index, 0, tgt_2).view(-1)

        def get_angle(vec1, vec2):
            vec1 = vec1 / (torch.norm(vec1, keepdim=True, dim=-1) + 1e-6)
            vec2 = vec2 / (torch.norm(vec2, keepdim=True, dim=-1) + 1e-6)
            return torch.einsum("nc,nc->n", vec1, vec2)

        true_angle = get_angle(
            pos_y[tgt_1] - pos_y[nei_src_index], pos_y[tgt_2] - pos_y[nei_src_index]
        )
        pred_angle = get_angle(
            pos_x[tgt_1] - pos_x[nei_src_index], pos_x[tgt_2] - pos_x[nei_src_index]
        )
        angle_loss = torch.mean(F.l1_loss(pred_angle, true_angle))

        return bond_loss, angle_loss
//  [python]
def aux_loss(pos_y, pos_x, batch):
        edge_index = batch.edge_index
        src = edge_index[0]
        tgt = edge_index[1]
        true_bond = torch.linalg.norm(pos_y[src] - pos_y[tgt], dim=-1)
        pred_bond = torch.linalg.norm(pos_x[src] - pos_x[tgt], dim=-1)
        bond_loss = torch.mean(F.l1_loss(pred_bond, true_bond))

        nei_src_index = batch.nei_src_index.view(-1)
        nei_tgt_index = batch.nei_tgt_index
        nei_tgt_mask = batch.nei_tgt_mask

        random_tgt_index = pos_y.new_zeros(nei_tgt_index.size()).uniform_()
        random_tgt_index = torch.where(
            nei_tgt_mask, pos_y.new_zeros(nei_tgt_index.size()), random_tgt_index
        )
        random_tgt_index_sort = torch.sort(random_tgt_index, descending=True, dim=0)[1][:2]
        tgt_1, tgt_2 = random_tgt_index_sort[0].unsqueeze(0), random_tgt_index_sort[1].unsqueeze(0)

        tgt_1 = torch.gather(nei_tgt_index, 0, tgt_1).view(-1)
        tgt_2 = torch.gather(nei_tgt_index, 0, tgt_2).view(-1)

        def get_angle(vec1, vec2):
            vec1 = vec1 / (torch.linalg.norm(vec1, keepdim=True, dim=-1) + 1e-6)
            vec2 = vec2 / (torch.linalg.norm(vec2, keepdim=True, dim=-1) + 1e-6)
            return torch.einsum("nc,nc->n", vec1, vec2)

        true_angle = get_angle(
            pos_y[tgt_1] - pos_y[nei_src_index], pos_y[tgt_2] - pos_y[nei_src_index]
        )
        pred_angle = get_angle(
            pos_x[tgt_1] - pos_x[nei_src_index], pos_x[tgt_2] - pos_x[nei_src_index]
        )
        angle_loss = torch.mean(F.l1_loss(pred_angle, true_angle))

        return bond_loss, angle_loss
//  [python]
def get_limb_lens(x):
    '''
        Input: (N, T, 17, 3)
        Output: (N, T, 16)
    '''
    limbs_id = [[0,1], [1,2], [2,3],
         [0,4], [4,5], [5,6],
         [0,7], [7,8], [8,9], [9,10],
         [8,11], [11,12], [12,13],
         [8,14], [14,15], [15,16]
        ]
    limbs = x[:,:,limbs_id,:]
    limbs = limbs[:,:,:,0,:]-limbs[:,:,:,1,:]
    limb_lens = torch.norm(limbs, dim=-1)
    return limb_lens
//  [python]
def get_limb_lens(x):
    '''
        Input: (N, T, 17, 3)
        Output: (N, T, 16)
    '''
    limbs_id = [[0,1], [1,2], [2,3],
         [0,4], [4,5], [5,6],
         [0,7], [7,8], [8,9], [9,10],
         [8,11], [11,12], [12,13],
         [8,14], [14,15], [15,16]
        ]
    limbs = x[:,:,limbs_id,:]
    limbs = limbs[:,:,:,0,:]-limbs[:,:,:,1,:]
    limb_lens = torch.linalg.norm(limbs, dim=-1)
    return limb_lens
//  [python]
def clip(inputs, ori_weight):

    with torch.no_grad():
        diff = inputs - ori_weight
        norm = torch.norm(diff, dim=-1)  #(B, 3)
        scale_factor = 3.5 / (norm + 1e-9)
        scale_factor = torch.clamp(scale_factor, max=1.)  # [B, 3]
        diff = diff * scale_factor[:, :, None] 

    return ori_weight + diff
//  [python]
def clip(inputs, ori_weight):

    with torch.no_grad():
        diff = inputs - ori_weight
        norm = torch.linalg.norm(diff, dim=-1)  #(B, 3)
        scale_factor = 3.5 / (norm + 1e-9)
        scale_factor = torch.clamp(scale_factor, max=1.)  # [B, 3]
        diff = diff * scale_factor[:, :, None] 

    return ori_weight + diff
//  [python]
def update_weight_Boosting(self, weight_mat, dist_old, dist_new):
        epsilon = 1e-5
        dist_old = dist_old.detach()
        dist_new = dist_new.detach()
        ind = dist_new > dist_old + epsilon
        weight_mat[ind] = weight_mat[ind] * (1 + torch.sigmoid(dist_new[ind] - dist_old[ind]))
        weight_norm = torch.norm(weight_mat, dim=1, p=1)
        weight_mat = weight_mat / weight_norm.t().unsqueeze(1).repeat(1, self.len_seq)
        return weight_mat
//  [python]
def update_weight_Boosting(self, weight_mat, dist_old, dist_new):
        epsilon = 1e-5
        dist_old = dist_old.detach()
        dist_new = dist_new.detach()
        ind = dist_new > dist_old + epsilon
        weight_mat[ind] = weight_mat[ind] * (1 + torch.sigmoid(dist_new[ind] - dist_old[ind]))
        weight_norm = torch.linalg.norm(weight_mat, dim=1, ord=1)
        weight_mat = weight_mat / weight_norm.t().unsqueeze(1).repeat(1, self.len_seq)
        return weight_mat
//  [python]
def embeddings_to_cosine_similarity(E, sigma=1.0):
    '''
    Build a pairwise symmetrical cosine similarity matrix
    diganal is set as zero
    '''

    dot = torch.abs_(torch.mm(E, E.t())) # E[i]E[j]
    norm = torch.norm(E, 2, 1) # ||E[i]||
    x = torch.div(dot, norm) # E[i]E[j]/||E[j]||
    x = torch.div(x, torch.unsqueeze(norm, 0)) # E[i]E[j]/(||E[j]||*||E[i]||)
    x = x.div_(sigma)

    return torch.max(x, x.t()).fill_diagonal_(0)
//  [python]
def embeddings_to_cosine_similarity(E, sigma=1.0):
    '''
    Build a pairwise symmetrical cosine similarity matrix
    diganal is set as zero
    '''

    dot = torch.abs_(torch.mm(E, E.t())) # E[i]E[j]
    norm = torch.linalg.norm(E, 2, 1) # ||E[i]||
    x = torch.div(dot, norm) # E[i]E[j]/||E[j]||
    x = torch.div(x, torch.unsqueeze(norm, 0)) # E[i]E[j]/(||E[j]||*||E[i]||)
    x = x.div_(sigma)

    return torch.max(x, x.t()).fill_diagonal_(0)
//  [python]
def FISTA(X, Wd, alpha, tolerance, max_steps=-1, debug=False):
    """
    Dimension notation:
        B - Number of samples. Usually number of patches in image times batch size
        K - Number of atoms in dictionary
        d - Dimensionality of atoms in dictionary


    :param X: Input - Signal to find sparse coding against. Dimenions: d x B
    :param Wd: Dictionary - Tensor of atoms we want to get a sparse linear combination of. Dimensions: d x K
    :param alpha:  Float. Sparsity weight
    :param tolerance: Float. Stop optimization once the improvement drops below this term.
    :param max_steps: Int. Stops optimizaiton after max_steps iterations. Ignored if < 0.
    :param debug: Bool. Output debug message if True.

    :return: Z: linear coefficients for Sparse Code solution. Dimensions: K x B
    """
    if debug:
        print('Doing FISTA with alpha={}'.format(alpha))
    L = get_largest_eigenvalue(Wd.t().mm(Wd))
    prev_Z = torch.zeros((Wd.size()[1], X.size()[1]), dtype=X.dtype).to(Wd.device)
    Y = prev_Z
    momentum = torch.ones(1).to(Wd.device)

    max_distance = 2 * tolerance
    step = 0
    while max_distance > tolerance or (step == 1 and torch.isnan(max_distance).item()):
        # Update sparse code
        cur_Z = fista_step(L, Wd, X, alpha, Y)
        Y, momentum = fista_momentum(cur_Z, prev_Z, momentum)
        # next_Z = cur_Z

        # calculate distance for tolerance
        distance_nominator = torch.norm(Wd.mm(prev_Z - cur_Z), p=2, dim=0)
        distance_denominator = torch.norm(Wd.mm(prev_Z),p=2, dim=0)
        distance = distance_nominator / (distance_denominator + np.finfo(float).eps)
        max_distance = distance.max()

        prev_Z = cur_Z

        step += 1
        if debug and step % 1 == 0:
            print('Step {}, code improvement: {}, below tolerance: {}.'.format(step, max_distance, (distance < tolerance).float().mean().item()))

        if step >= max_steps > 0:
            break

    return cur_Z, step
//  [python]
def FISTA(X, Wd, alpha, tolerance, max_steps=-1, debug=False):
    """
    Dimension notation:
        B - Number of samples. Usually number of patches in image times batch size
        K - Number of atoms in dictionary
        d - Dimensionality of atoms in dictionary


    :param X: Input - Signal to find sparse coding against. Dimenions: d x B
    :param Wd: Dictionary - Tensor of atoms we want to get a sparse linear combination of. Dimensions: d x K
    :param alpha:  Float. Sparsity weight
    :param tolerance: Float. Stop optimization once the improvement drops below this term.
    :param max_steps: Int. Stops optimizaiton after max_steps iterations. Ignored if < 0.
    :param debug: Bool. Output debug message if True.

    :return: Z: linear coefficients for Sparse Code solution. Dimensions: K x B
    """
    if debug:
        print('Doing FISTA with alpha={}'.format(alpha))
    L = get_largest_eigenvalue(Wd.t().mm(Wd))
    prev_Z = torch.zeros((Wd.size()[1], X.size()[1]), dtype=X.dtype).to(Wd.device)
    Y = prev_Z
    momentum = torch.ones(1).to(Wd.device)

    max_distance = 2 * tolerance
    step = 0
    while max_distance > tolerance or (step == 1 and torch.isnan(max_distance).item()):
        # Update sparse code
        cur_Z = fista_step(L, Wd, X, alpha, Y)
        Y, momentum = fista_momentum(cur_Z, prev_Z, momentum)
        # next_Z = cur_Z

        # calculate distance for tolerance
        distance_nominator = torch.linalg.norm(Wd.mm(prev_Z - cur_Z), ord=2, dim=0)
        distance_denominator = torch.linalg.norm(Wd.mm(prev_Z),ord=2, dim=0)
        distance = distance_nominator / (distance_denominator + np.finfo(float).eps)
        max_distance = distance.max()

        prev_Z = cur_Z

        step += 1
        if debug and step % 1 == 0:
            print('Step {}, code improvement: {}, below tolerance: {}.'.format(step, max_distance, (distance < tolerance).float().mean().item()))

        if step >= max_steps > 0:
            break

    return cur_Z, step
//  [python]
def dictionary_learning():


    n = 30

    np.random.seed(1)
    m = 10*n**2   # sample complexity
    theta = 0.3   # sparsity level
    Y = norm.ppf(np.random.rand(n,m)) * (norm.ppf(np.random.rand(n,m)) <= theta)  # Bernoulli-Gaussian model
    Y = torch.from_numpy(Y).to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"q": [n,1]}


    # def comb_fn(X_struct):
    #     q = X_struct.q

    #     # objective function
    #     qtY = q.T @ Y
    #     f = 1/m * torch.norm(qtY, p = 1)

    #     # inequality constraint, matrix form
    #     ci = None

    #     # equality constraint
    #     ce = pygransoStruct()
    #     ce.c1 = q.T @ q - 1

    #     return [f,ci,ce]

    # Without AD
    def comb_fn(X_struct):
        q = X_struct.q
        
        # objective function
        qtY = q.T @ Y
        f = 1/m * torch.norm(qtY, p = 1).item()
        f_grad = 1/m*Y@torch.sign(Y.T@q)

        # inequality constraint, matrix form
        ci = None
        ci_grad = None

        # equality constraint 
        ce = q.T @ q - 1
        ce_grad = 2*q

        return [f,f_grad,ci,ci_grad,ce,ce_grad]

    opts = pygransoStruct()
    opts.QPsolver = 'osqp'
    opts.maxit = 20
    np.random.seed(1)
    x0 = norm.ppf(np.random.rand(n,1))
    x0 /= la.norm(x0,2)
    opts.x0 = torch.from_numpy(x0).to(device=device, dtype=torch_dtype)
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device


    opts.print_frequency = 10
    opts.globalAD = False

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 3/7 passed (Dictionary Learning)")
//  [python]
def dictionary_learning():


    n = 30

    np.random.seed(1)
    m = 10*n**2   # sample complexity
    theta = 0.3   # sparsity level
    Y = norm.ppf(np.random.rand(n,m)) * (norm.ppf(np.random.rand(n,m)) <= theta)  # Bernoulli-Gaussian model
    Y = torch.from_numpy(Y).to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"q": [n,1]}


    # def comb_fn(X_struct):
    #     q = X_struct.q

    #     # objective function
    #     qtY = q.T @ Y
    #     f = 1/m * torch.linalg.norm(qtY, p = 1)

    #     # inequality constraint, matrix form
    #     ci = None

    #     # equality constraint
    #     ce = pygransoStruct()
    #     ce.c1 = q.T @ q - 1

    #     return [f,ci,ce]

    # Without AD
    def comb_fn(X_struct):
        q = X_struct.q
        
        # objective function
        qtY = q.T @ Y
        f = 1/m * torch.linalg.norm(qtY, p = 1).item()
        f_grad = 1/m*Y@torch.sign(Y.T@q)

        # inequality constraint, matrix form
        ci = None
        ci_grad = None

        # equality constraint 
        ce = q.T @ q - 1
        ce_grad = 2*q

        return [f,f_grad,ci,ci_grad,ce,ce_grad]

    opts = pygransoStruct()
    opts.QPsolver = 'osqp'
    opts.maxit = 20
    np.random.seed(1)
    x0 = norm.ppf(np.random.rand(n,1))
    x0 /= la.norm(x0,2)
    opts.x0 = torch.from_numpy(x0).to(device=device, dtype=torch_dtype)
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device


    opts.print_frequency = 10
    opts.globalAD = False

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 3/7 passed (Dictionary Learning)")
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    Convert rotations given as axis/angle to quaternions.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    Convert rotations given as axis/angle to quaternions.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def compute(self, X, Y):
        """Compute adaptation loss

        Arguments:
            X {tensor} -- source matrix
            Y {tensor} -- target matrix

        Returns:
            [tensor] -- transfer loss
        """
        if self.loss_type in ("mmd_lin", "mmd"):
            mmdloss = MMD_loss(kernel_type="linear")
            loss = mmdloss(X, Y)
        elif self.loss_type == "coral":
            loss = CORAL(X, Y, self.device)
        elif self.loss_type in ("cosine", "cos"):
            loss = 1 - cosine(X, Y)
        elif self.loss_type == "kl":
            loss = kl_div(X, Y)
        elif self.loss_type == "js":
            loss = js(X, Y)
        elif self.loss_type == "mine":
            mine_model = Mine_estimator(input_dim=self.input_dim, hidden_dim=60).to(self.device)
            loss = mine_model(X, Y)
        elif self.loss_type == "adv":
            loss = adv(X, Y, self.device, input_dim=self.input_dim, hidden_dim=32)
        elif self.loss_type == "mmd_rbf":
            mmdloss = MMD_loss(kernel_type="rbf")
            loss = mmdloss(X, Y)
        elif self.loss_type == "pairwise":
            pair_mat = pairwise_dist(X, Y)
            loss = torch.norm(pair_mat)

        return loss
//  [python]
def compute(self, X, Y):
        """Compute adaptation loss

        Arguments:
            X {tensor} -- source matrix
            Y {tensor} -- target matrix

        Returns:
            [tensor] -- transfer loss
        """
        if self.loss_type in ("mmd_lin", "mmd"):
            mmdloss = MMD_loss(kernel_type="linear")
            loss = mmdloss(X, Y)
        elif self.loss_type == "coral":
            loss = CORAL(X, Y, self.device)
        elif self.loss_type in ("cosine", "cos"):
            loss = 1 - cosine(X, Y)
        elif self.loss_type == "kl":
            loss = kl_div(X, Y)
        elif self.loss_type == "js":
            loss = js(X, Y)
        elif self.loss_type == "mine":
            mine_model = Mine_estimator(input_dim=self.input_dim, hidden_dim=60).to(self.device)
            loss = mine_model(X, Y)
        elif self.loss_type == "adv":
            loss = adv(X, Y, self.device, input_dim=self.input_dim, hidden_dim=32)
        elif self.loss_type == "mmd_rbf":
            mmdloss = MMD_loss(kernel_type="rbf")
            loss = mmdloss(X, Y)
        elif self.loss_type == "pairwise":
            pair_mat = pairwise_dist(X, Y)
            loss = torch.linalg.norm(pair_mat)

        return loss
//  [python]
def quaternion_to_axis_angle(quaternions):
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def quaternion_to_axis_angle(quaternions):
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.linalg.norm(quaternions[..., 1:], ord=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def pairwise_cosine_similarity(E, C, temp=1.0):
    '''
    Build a pairwise cosine similarity matrix.
    '''
    dot = torch.abs_(torch.mm(E, C.t())) # E[i]C[j]
    # dot = torch.mm(E, C.t()) # E[i]C[j]
    # print(dot.size())
    E = torch.norm(E, 2, 1) # ||E[j]||
    # print(E.size(), norm_E.size())
    C = torch.norm(C, 2, 1) # ||C[i]||
    # print(C.size(), norm_C.size())
    x = torch.div(dot, E.view(-1, 1)) # E[i]E[j]/||E[j]||
    x = torch.div(x, C.view(1, -1)) # E[i]E[j]/(||E[j]||*||E[i]||)
    return x.div_(temp)
//  [python]
def pairwise_cosine_similarity(E, C, temp=1.0):
    '''
    Build a pairwise cosine similarity matrix.
    '''
    dot = torch.abs_(torch.mm(E, C.t())) # E[i]C[j]
    # dot = torch.mm(E, C.t()) # E[i]C[j]
    # print(dot.size())
    E = torch.linalg.norm(E, 2, 1) # ||E[j]||
    # print(E.size(), norm_E.size())
    C = torch.linalg.norm(C, 2, 1) # ||C[i]||
    # print(C.size(), norm_C.size())
    x = torch.div(dot, E.view(-1, 1)) # E[i]E[j]/||E[j]||
    x = torch.div(x, C.view(1, -1)) # E[i]E[j]/(||E[j]||*||E[i]||)
    return x.div_(temp)
//  [python]
def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    return total_norm
//  [python]
def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.linalg.norm(torch.stack([torch.linalg.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    return total_norm
//  [python]
def gram_schmidt(V):
    n,N=V.shape #dimension, cardinal
    W=th.zeros_like(V)
    R=th.zeros((N,N)).double().to(V.device)
    W[:,0]=V[:,0]/th.norm(V[:,0])
    R[0,0]=W[:,0].dot(V[:,0])
    for i in range(1,N):
        proj=th.zeros(n).double().to(V.device)
        for j in range(i):
            proj=proj+V[:,i].dot(W[:,j])*W[:,j]
            R[j,i]=W[:,j].dot(V[:,i])
        if(isclose(th.norm(V[:,i]-proj),th.DoubleTensor([0]).to(V.device))):
            W[:,i]=V[:,i]/th.norm(V[:,i])
        else:
            W[:,i]=(V[:,i]-proj)/th.norm(V[:,i]-proj)
        R[i,i]=W[:,i].dot(V[:,i])
    return W,R
//  [python]
def gram_schmidt(V):
    n,N=V.shape #dimension, cardinal
    W=th.zeros_like(V)
    R=th.zeros((N,N)).double().to(V.device)
    W[:,0]=V[:,0]/th.linalg.norm(V[:,0])
    R[0,0]=W[:,0].dot(V[:,0])
    for i in range(1,N):
        proj=th.zeros(n).double().to(V.device)
        for j in range(i):
            proj=proj+V[:,i].dot(W[:,j])*W[:,j]
            R[j,i]=W[:,j].dot(V[:,i])
        if(isclose(th.linalg.norm(V[:,i]-proj),th.DoubleTensor([0]).to(V.device))):
            W[:,i]=V[:,i]/th.linalg.norm(V[:,i])
        else:
            W[:,i]=(V[:,i]-proj)/th.linalg.norm(V[:,i]-proj)
        R[i,i]=W[:,i].dot(V[:,i])
    return W,R
//  [python]
def get_rays_tourism(H, W, kinv, pose):
    """
    phototourism camera intrinsics are defined by H, W and kinv.
    Args:
        H: image height
        W: image width
        kinv (3, 3): inverse of camera intrinsic
        pose (4, 4): camera extrinsic
    Returns:
        rays_o (H, W, 3): ray origins
        rays_d (H, W, 3): ray directions
    """
    yy, xx = torch.meshgrid(torch.arange(0., H, device=kinv.device),
                            torch.arange(0., W, device=kinv.device),
                            indexing='ij')
    pixco = torch.stack([xx, yy, torch.ones_like(xx)], dim=-1)

    directions = torch.matmul(pixco, kinv.T)  # (H, W, 3)

    rays_d = torch.matmul(directions, pose[:3, :3].T)
    rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)  # (H, W, 3)

    rays_o = pose[:3, -1].expand_as(rays_d)  # (H, W, 3)

    return rays_o, rays_d
//  [python]
def get_rays_tourism(H, W, kinv, pose):
    """
    phototourism camera intrinsics are defined by H, W and kinv.
    Args:
        H: image height
        W: image width
        kinv (3, 3): inverse of camera intrinsic
        pose (4, 4): camera extrinsic
    Returns:
        rays_o (H, W, 3): ray origins
        rays_d (H, W, 3): ray directions
    """
    yy, xx = torch.meshgrid(torch.arange(0., H, device=kinv.device),
                            torch.arange(0., W, device=kinv.device),
                            indexing='ij')
    pixco = torch.stack([xx, yy, torch.ones_like(xx)], dim=-1)

    directions = torch.matmul(pixco, kinv.T)  # (H, W, 3)

    rays_d = torch.matmul(directions, pose[:3, :3].T)
    rays_d = rays_d / torch.linalg.norm(rays_d, dim=-1, keepdim=True)  # (H, W, 3)

    rays_o = pose[:3, -1].expand_as(rays_d)  # (H, W, 3)

    return rays_o, rays_d
//  [python]
def obstacle_cost_gmm(xt):

    Bs, D = xt.shape[:-1], xt.shape[-1]
    assert D == 2
    xt = xt.reshape(-1, xt.shape[-1])

    batch_xt = xt.shape[0]

    centers, radius = obstacle_cfg_gmm()

    obs1 = torch.tensor(centers[0]).repeat((batch_xt, 1)).to(xt.device)
    obs2 = torch.tensor(centers[1]).repeat((batch_xt, 1)).to(xt.device)
    obs3 = torch.tensor(centers[2]).repeat((batch_xt, 1)).to(xt.device)

    dist1 = torch.norm(xt - obs1, dim=-1)
    dist2 = torch.norm(xt - obs2, dim=-1)
    dist3 = torch.norm(xt - obs3, dim=-1)

    cost1 = F.softplus(100 * (radius - dist1), beta=1, threshold=20)
    cost2 = F.softplus(100 * (radius - dist2), beta=1, threshold=20)
    cost3 = F.softplus(100 * (radius - dist3), beta=1, threshold=20)
    return (cost1 + cost2 + cost3).reshape(*Bs)
//  [python]
def obstacle_cost_gmm(xt):

    Bs, D = xt.shape[:-1], xt.shape[-1]
    assert D == 2
    xt = xt.reshape(-1, xt.shape[-1])

    batch_xt = xt.shape[0]

    centers, radius = obstacle_cfg_gmm()

    obs1 = torch.tensor(centers[0]).repeat((batch_xt, 1)).to(xt.device)
    obs2 = torch.tensor(centers[1]).repeat((batch_xt, 1)).to(xt.device)
    obs3 = torch.tensor(centers[2]).repeat((batch_xt, 1)).to(xt.device)

    dist1 = torch.linalg.norm(xt - obs1, dim=-1)
    dist2 = torch.linalg.norm(xt - obs2, dim=-1)
    dist3 = torch.linalg.norm(xt - obs3, dim=-1)

    cost1 = F.softplus(100 * (radius - dist1), beta=1, threshold=20)
    cost2 = F.softplus(100 * (radius - dist2), beta=1, threshold=20)
    cost3 = F.softplus(100 * (radius - dist3), beta=1, threshold=20)
    return (cost1 + cost2 + cost3).reshape(*Bs)
//  [python]
def fwd_solver(f, z_init, threshold, max_iter, mode='abs'):
    # ipdb.set_trace()
    z_prev, z = z_init, f(z_init)
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        abs_diff = torch.norm(z_prev - z).item()
        rel_diff = abs_diff / (torch.norm(z_prev).item() + 1e-9)
        diff_dict = {'abs': abs_diff, 'rel': rel_diff}
        # print(f'abs_diff: {diff_dict["abs"]}, rel_diff: {rel_diff}')
        if diff_dict[mode] < threshold:
            # print(nstep, abs_diff)
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def fwd_solver(f, z_init, threshold, max_iter, mode='abs'):
    # ipdb.set_trace()
    z_prev, z = z_init, f(z_init)
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        abs_diff = torch.linalg.norm(z_prev - z).item()
        rel_diff = abs_diff / (torch.linalg.norm(z_prev).item() + 1e-9)
        diff_dict = {'abs': abs_diff, 'rel': rel_diff}
        # print(f'abs_diff: {diff_dict["abs"]}, rel_diff: {rel_diff}')
        if diff_dict[mode] < threshold:
            # print(nstep, abs_diff)
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def update_metric(self):
        r"""
        As soon as the model has seen all the data points (i.e. at the end of 1 loop)
        we update the final metric function using \mu(x_i) as centroids
        """
        # convert to 1 big tensor
        self.M_tens = torch.cat(self.M)
        self.centroids_tens = torch.cat(self.centroids)

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.temperature ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(z.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.temperature ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(z.device)

        self.G = G
        self.G_inv = G_inv
        self.M = []
        self.centroids = []
//  [python]
def update_metric(self):
        r"""
        As soon as the model has seen all the data points (i.e. at the end of 1 loop)
        we update the final metric function using \mu(x_i) as centroids
        """
        # convert to 1 big tensor
        self.M_tens = torch.cat(self.M)
        self.centroids_tens = torch.cat(self.centroids)

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.linalg.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.temperature ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(z.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.linalg.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.temperature ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(z.device)

        self.G = G
        self.G_inv = G_inv
        self.M = []
        self.centroids = []
//  [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]
    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):
        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]
    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):
        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 new_solver(q, h, vjp_func, max_iter):
    nstep = 0
    # q and h should be initialized as v.
    while nstep < max_iter:
        # func(h) = autograd.grad(new_Z, Z, h, retain_graph=True)
        h = vjp_func(h)
        q = q + h
        print('current step: ', nstep, 'torch.norm(h).item(): ', torch.norm(h).item())
        nstep += 1
    return q
//  [python]
def new_solver(q, h, vjp_func, max_iter):
    nstep = 0
    # q and h should be initialized as v.
    while nstep < max_iter:
        # func(h) = autograd.grad(new_Z, Z, h, retain_graph=True)
        h = vjp_func(h)
        q = q + h
        print('current step: ', nstep, 'torch.linalg.norm(h).item(): ', torch.linalg.norm(h).item())
        nstep += 1
    return q
//  [python]
def analyse_fwd_solver(f, z_init, threshold, max_iter):
    z_prev, z = z_init, f(z_init)
    z_list = [z_prev, z]
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        z_list.append(z)
        abs_diff = torch.norm(z_prev - z).item()
        if abs_diff < threshold:
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def analyse_fwd_solver(f, z_init, threshold, max_iter):
    z_prev, z = z_init, f(z_init)
    z_list = [z_prev, z]
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        z_list.append(z)
        abs_diff = torch.linalg.norm(z_prev - z).item()
        if abs_diff < threshold:
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def _safe_norm(v):
    if not torch.isfinite(v).all():
        return np.inf
    return torch.norm(v)
//  [python]
def _safe_norm(v):
    if not torch.isfinite(v).all():
        return np.inf
    return torch.linalg.norm(v)
//  [python]
def abs(self, x, y):
        num_examples = x.size()[0]

        #Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h**(self.d/self.p))*torch.norm(x.view(num_examples,-1) - y.view(num_examples,-1), self.p, 1)

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms
//  [python]
def abs(self, x, y):
        num_examples = x.size()[0]

        #Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h**(self.d/self.p))*torch.linalg.norm(x.view(num_examples,-1) - y.view(num_examples,-1), ord=self.p, dim=1)

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms
//  [python]
def get_grad_norm(model):
    parameters = [p for p in model.parameters() if p.grad is not None and p.requires_grad]
    if len(parameters) == 0:
        total_norm = 0.0
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), 2.0) for p in parameters]), 2.0).item()
    return total_norm
//  [python]
def get_grad_norm(model):
    parameters = [p for p in model.parameters() if p.grad is not None and p.requires_grad]
    if len(parameters) == 0:
        total_norm = 0.0
    else:
        total_norm = torch.linalg.norm(torch.stack([torch.linalg.norm(p.grad.detach(), 2.0) for p in parameters]), 2.0).item()
    return total_norm