jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def balance_norm(factors):
	norm = 1
	for f in factors:
		norm = norm * torch.norm(f, dim=0)
	even_split = (norm + 1e-15)#.pow(1.0 / len(factors))
	for i in range(len(factors)):
		factors[i] = factors[i] / (torch.norm(factors[i], dim=0) + 1e-15) #* even_split
	factors[-1] = factors[-1] * even_split
//  [python]
def balance_norm(factors):
	norm = 1
	for f in factors:
		norm = norm * torch.linalg.norm(f, dim=0)
	even_split = (norm + 1e-15)#.pow(1.0 / len(factors))
	for i in range(len(factors)):
		factors[i] = factors[i] / (torch.linalg.norm(factors[i], dim=0) + 1e-15) #* even_split
	factors[-1] = factors[-1] * even_split
//  [python]
def regress_meta(self, replay_iter, step):
        if self.solved_meta is not None:
            return self.solved_meta
        feats, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            data = next(replay_iter)
            data = self.wm.preprocess(data)
            embed = self.wm.encoder(data)
            post, prior = self.wm.rssm.observe(
                embed, data['action'], data['is_first'])
            feat = self.wm.rssm.get_feat(post)
            B, T, _ = feat.shape

            feats.append(feat.reshape(B*T,-1))
            reward.append(data['reward'].reshape(B*T,-1))
            batch_size += B*T
        feats, reward = torch.cat(feats, 0), torch.cat(reward, 0)

        rep = self.aps(feats)
        skill = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        skill = skill / torch.norm(skill)
        skill = skill.cpu().numpy()
        meta = OrderedDict()
        meta['skill'] = skill

        print("skill selected : ", meta['skill'])

        # save for evaluation
        self.solved_meta = meta
        self._task_behavior.solved_meta = meta 
        return meta
//  [python]
def regress_meta(self, replay_iter, step):
        if self.solved_meta is not None:
            return self.solved_meta
        feats, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            data = next(replay_iter)
            data = self.wm.preprocess(data)
            embed = self.wm.encoder(data)
            post, prior = self.wm.rssm.observe(
                embed, data['action'], data['is_first'])
            feat = self.wm.rssm.get_feat(post)
            B, T, _ = feat.shape

            feats.append(feat.reshape(B*T,-1))
            reward.append(data['reward'].reshape(B*T,-1))
            batch_size += B*T
        feats, reward = torch.cat(feats, 0), torch.cat(reward, 0)

        rep = self.aps(feats)
        skill = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        skill = skill / torch.linalg.norm(skill)
        skill = skill.cpu().numpy()
        meta = OrderedDict()
        meta['skill'] = skill

        print("skill selected : ", meta['skill'])

        # save for evaluation
        self.solved_meta = meta
        self._task_behavior.solved_meta = meta 
        return meta
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=1000, min_delta=1e-9):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf

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

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

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

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

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

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=1000, min_delta=1e-9):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf

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

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

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

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

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

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def inner_loss(B_bar, w_bar, X_list, Y_list, lambda_1, lambda_2, B, w, A = None):
    K = len(w)
    loss_1 = 0
    for ind_sample in range(len(X_list)):
        X_l = X_list[ind_sample]
        y_l = Y_list[ind_sample]
        pred_l = 0
        for k in range(K):
            if A is None:
                pred_l += B[k]@Herm(B[k])@X_l@w[k]
            else:
                pred_l += A[k]@X_l@w[k]
        loss_1 += torch.norm(y_l-pred_l)**2
    loss_2 = 0
    if lambda_1 == 0:
        pass
    else:
        for k in range(K):
            if A is None:
                loss_2 += lambda_1 * torch.norm(Herm(B_bar[k])- Herm(B_bar[k])@B[k]@Herm(B[k]) )**2
            else:
                loss_2 += lambda_1 * torch.norm(Herm(B_bar[k])- Herm(B_bar[k])@A[k] )**2
    loss_3 = 0
    if lambda_2 == 0:
        pass
    else:
        for k in range(K):
            loss_3 += lambda_2 * torch.norm(w[k] - w_bar[k])**2 # it was w!!
    return loss_1 + loss_2 + loss_3
//  [python]
def inner_loss(B_bar, w_bar, X_list, Y_list, lambda_1, lambda_2, B, w, A = None):
    K = len(w)
    loss_1 = 0
    for ind_sample in range(len(X_list)):
        X_l = X_list[ind_sample]
        y_l = Y_list[ind_sample]
        pred_l = 0
        for k in range(K):
            if A is None:
                pred_l += B[k]@Herm(B[k])@X_l@w[k]
            else:
                pred_l += A[k]@X_l@w[k]
        loss_1 += torch.linalg.norm(y_l-pred_l)**2
    loss_2 = 0
    if lambda_1 == 0:
        pass
    else:
        for k in range(K):
            if A is None:
                loss_2 += lambda_1 * torch.linalg.norm(Herm(B_bar[k])- Herm(B_bar[k])@B[k]@Herm(B[k]) )**2
            else:
                loss_2 += lambda_1 * torch.linalg.norm(Herm(B_bar[k])- Herm(B_bar[k])@A[k] )**2
    loss_3 = 0
    if lambda_2 == 0:
        pass
    else:
        for k in range(K):
            loss_3 += lambda_2 * torch.linalg.norm(w[k] - w_bar[k])**2 # it was w!!
    return loss_1 + loss_2 + loss_3
//  [python]
def pairwise_distance_mse(U, V, D, reg=1):
    d_uv = torch.cdist(U, V)
    l = torch.norm(D - d_uv)**2 / D.numel() + reg * (torch.norm(U) **
                                                     2 / U.numel() + torch.norm(V)**2 / V.numel())  # MSE per entry
    return l
//  [python]
def pairwise_distance_mse(U, V, D, reg=1):
    d_uv = torch.cdist(U, V)
    l = torch.linalg.norm(D - d_uv)**2 / D.numel() + reg * (torch.linalg.norm(U)**2 / U.numel() + torch.linalg.norm(V)**2 / V.numel())  # MSE per entry
    return l
//  [python]
def match_features(feat_src, feat_ref, metric='l2'):
    """ Compute pairwise distance between features

    Args:
        feat_src: (B, J, C)
        feat_ref: (B, K, C)
        metric: either 'angle' or 'l2' (squared euclidean)

    Returns:
        Matching matrix (B, J, K). i'th row describes how well the i'th point
         in the src agrees with every point in the ref.
    """
    assert feat_src.shape[-1] == feat_ref.shape[-1]

    if metric == 'l2':
        dist_matrix = square_distance(feat_src, feat_ref)
    elif metric == 'angle':
        feat_src_norm = feat_src / (torch.norm(feat_src, dim=-1, keepdim=True) + _EPS)
        feat_ref_norm = feat_ref / (torch.norm(feat_ref, dim=-1, keepdim=True) + _EPS)

        dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
    else:
        raise NotImplementedError

    return dist_matrix
//  [python]
def match_features(feat_src, feat_ref, metric='l2'):
    """ Compute pairwise distance between features

    Args:
        feat_src: (B, J, C)
        feat_ref: (B, K, C)
        metric: either 'angle' or 'l2' (squared euclidean)

    Returns:
        Matching matrix (B, J, K). i'th row describes how well the i'th point
         in the src agrees with every point in the ref.
    """
    assert feat_src.shape[-1] == feat_ref.shape[-1]

    if metric == 'l2':
        dist_matrix = square_distance(feat_src, feat_ref)
    elif metric == 'angle':
        feat_src_norm = feat_src / (torch.linalg.norm(feat_src, dim=-1, keepdim=True) + _EPS)
        feat_ref_norm = feat_ref / (torch.linalg.norm(feat_ref, dim=-1, keepdim=True) + _EPS)

        dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
    else:
        raise NotImplementedError

    return dist_matrix
//  [python]
def _test_vision_model(self, float_model):
        float_model.to('cpu')
        float_model.eval()
        float_model.fuse_model()
        float_model.qconfig = torch.ao.quantization.default_qconfig
        img_data = [(torch.rand(2, 3, 224, 224, dtype=torch.float), torch.randint(0, 1, (2,), dtype=torch.long)) for _ in range(2)]
        qmodel = quantize(float_model, torch.ao.quantization.default_eval_fn, [img_data], inplace=False)

        wt_compare_dict = compare_weights(float_model.state_dict(), qmodel.state_dict())

        def compute_error(x, y):
            Ps = torch.norm(x)
            Pn = torch.norm(x - y)
            return 20 * torch.log10(Ps / Pn)

        data = img_data[0][0]
        # Take in floating point and quantized model as well as input data, and returns a dict, with keys
        # corresponding to the quantized module names and each entry being a dictionary with two keys 'float' and
        # 'quantized', containing the activations of floating point and quantized model at matching locations.
        act_compare_dict = compare_model_outputs(float_model, qmodel, data)


        for key in act_compare_dict:
            compute_error(act_compare_dict[key]['float'][0], act_compare_dict[key]['quantized'][0].dequantize())

        prepare_model_outputs(float_model, qmodel)

        for data in img_data:
            float_model(data[0])
            qmodel(data[0])

        # Find the matching activation between floating point and quantized modules, and return a dict with key
        # corresponding to quantized module names and each entry being a dictionary with two keys 'float'
        # and 'quantized', containing the matching floating point and quantized activations logged by the logger
        act_compare_dict = get_matching_activations(float_model, qmodel)
//  [python]
def _test_vision_model(self, float_model):
        float_model.to('cpu')
        float_model.eval()
        float_model.fuse_model()
        float_model.qconfig = torch.ao.quantization.default_qconfig
        img_data = [(torch.rand(2, 3, 224, 224, dtype=torch.float), torch.randint(0, 1, (2,), dtype=torch.long)) for _ in range(2)]
        qmodel = quantize(float_model, torch.ao.quantization.default_eval_fn, [img_data], inplace=False)

        wt_compare_dict = compare_weights(float_model.state_dict(), qmodel.state_dict())

        def compute_error(x, y):
            Ps = torch.linalg.norm(x)
            Pn = torch.linalg.norm(x - y)
            return 20 * torch.log10(Ps / Pn)

        data = img_data[0][0]
        # Take in floating point and quantized model as well as input data, and returns a dict, with keys
        # corresponding to the quantized module names and each entry being a dictionary with two keys 'float' and
        # 'quantized', containing the activations of floating point and quantized model at matching locations.
        act_compare_dict = compare_model_outputs(float_model, qmodel, data)


        for key in act_compare_dict:
            compute_error(act_compare_dict[key]['float'][0], act_compare_dict[key]['quantized'][0].dequantize())

        prepare_model_outputs(float_model, qmodel)

        for data in img_data:
            float_model(data[0])
            qmodel(data[0])

        # Find the matching activation between floating point and quantized modules, and return a dict with key
        # corresponding to quantized module names and each entry being a dictionary with two keys 'float'
        # and 'quantized', containing the matching floating point and quantized activations logged by the logger
        act_compare_dict = get_matching_activations(float_model, qmodel)
//  [python]
def opt_multiple_w_gauss_seidel(B_bar, w_bar, X_list, Y_list, lambda_2, B, prev_w, max_iter=1, conv_criteria=CONV_CRITERIA):
    K = len(w_bar) # total number of different predictors
    R, P, Q = generate_P_Q_R_for_opt_w(X_list, Y_list, B, K)
    # init w
    if prev_w is None:
        raise NotImplementedError
        w = copy.deepcopy(w_bar)
    else:
        w = prev_w
    curr_loss = 999999
    for ind_iter in range(max_iter):
        for k in range(K): # coordinate descent
            R_sum_of_other_w = torch.zeros(w_bar[k].shape, dtype=torch.cdouble)
            for k_prime in range(K):
                if k == k_prime:
                    pass
                else:
                    R_sum_of_other_w += R[(k,k_prime)]@w[k_prime]
            tmp_1 = lambda_2*torch.eye(P[k].shape[0]) + P[k]
            tmp_2 = lambda_2*w_bar[k] + Q[k] - R_sum_of_other_w
            w[k] = torch.linalg.pinv(tmp_1)@tmp_2 # Guass-Seidel
        loss_tmp = inner_loss(B_bar, w_bar, X_list, Y_list, 0, lambda_2, B, w) #lambda_1 = 0 since we are fixing B here
        if loss_tmp<=conv_criteria:
            break 
        elif (loss_tmp <= curr_loss) and (torch.norm(curr_loss - loss_tmp) <= conv_criteria):
            break
        else:
            curr_loss = loss_tmp
    return w
//  [python]
def opt_multiple_w_gauss_seidel(B_bar, w_bar, X_list, Y_list, lambda_2, B, prev_w, max_iter=1, conv_criteria=CONV_CRITERIA):
    K = len(w_bar) # total number of different predictors
    R, P, Q = generate_P_Q_R_for_opt_w(X_list, Y_list, B, K)
    # init w
    if prev_w is None:
        raise NotImplementedError
        w = copy.deepcopy(w_bar)
    else:
        w = prev_w
    curr_loss = 999999
    for ind_iter in range(max_iter):
        for k in range(K): # coordinate descent
            R_sum_of_other_w = torch.zeros(w_bar[k].shape, dtype=torch.cdouble)
            for k_prime in range(K):
                if k == k_prime:
                    pass
                else:
                    R_sum_of_other_w += R[(k,k_prime)]@w[k_prime]
            tmp_1 = lambda_2*torch.eye(P[k].shape[0]) + P[k]
            tmp_2 = lambda_2*w_bar[k] + Q[k] - R_sum_of_other_w
            w[k] = torch.linalg.pinv(tmp_1)@tmp_2 # Guass-Seidel
        loss_tmp = inner_loss(B_bar, w_bar, X_list, Y_list, 0, lambda_2, B, w) #lambda_1 = 0 since we are fixing B here
        if loss_tmp<=conv_criteria:
            break 
        elif (loss_tmp <= curr_loss) and (torch.linalg.norm(curr_loss - loss_tmp) <= conv_criteria):
            break
        else:
            curr_loss = loss_tmp
    return w
//  [python]
def normalize_rotation(rotation, rotation_mode):
    if rotation_mode == 'quat_wxyz' or rotation_mode == 'quat_xyzw':
        rotation /= torch.norm(rotation, dim=-1, keepdim=True)
    elif rotation_mode == 'rot_matrix':
        rot_matrix = get_rot_matrix(rotation, rotation_mode)
        rotation[:, :3] = rot_matrix[:, :, 0]
        rotation[:, 3:6] = rot_matrix[:, :, 1]
    elif rotation_mode == 'euler_xyz_sx_cx':
        rot_sin_theta = rotation[:, :3]
        rot_cos_theta = rotation[:, 3:6]
        theta = torch.atan2(rot_sin_theta, rot_cos_theta)
        rotation[:, :3] = torch.sin(theta)
        rotation[:, 3:6] = torch.cos(theta)
    elif rotation_mode == 'euler_xyz':
        pass
    else:
        raise NotImplementedError
    return rotation
//  [python]
def normalize_rotation(rotation, rotation_mode):
    if rotation_mode == 'quat_wxyz' or rotation_mode == 'quat_xyzw':
        rotation /= torch.linalg.norm(rotation, dim=-1, keepdim=True)
    elif rotation_mode == 'rot_matrix':
        rot_matrix = get_rot_matrix(rotation, rotation_mode)
        rotation[:, :3] = rot_matrix[:, :, 0]
        rotation[:, 3:6] = rot_matrix[:, :, 1]
    elif rotation_mode == 'euler_xyz_sx_cx':
        rot_sin_theta = rotation[:, :3]
        rot_cos_theta = rotation[:, 3:6]
        theta = torch.atan2(rot_sin_theta, rot_cos_theta)
        rotation[:, :3] = torch.sin(theta)
        rotation[:, 3:6] = torch.cos(theta)
    elif rotation_mode == 'euler_xyz':
        pass
    else:
        raise NotImplementedError
    return rotation
//  [python]
def _cka(gram_x, gram_y, debiased=False):
    gram_x = center_gram(gram_x, unbiased=debiased)
    gram_y = center_gram(gram_y, unbiased=debiased)

    scaled_hsic = (gram_x.view(-1) * gram_y.view(-1)).sum()
    # normalization_x = torch.linalg.norm(gram_x)
    # normalization_y = torch.linalg.norm(gram_y)
    normalization_x = torch.norm(gram_x)
    normalization_y = torch.norm(gram_y)

    return scaled_hsic / (normalization_x * normalization_y)
//  [python]
def _cka(gram_x, gram_y, debiased=False):
    gram_x = center_gram(gram_x, unbiased=debiased)
    gram_y = center_gram(gram_y, unbiased=debiased)

    scaled_hsic = (gram_x.view(-1) * gram_y.view(-1)).sum()
    normalization_x = torch.linalg.norm(gram_x)
    normalization_y = torch.linalg.norm(gram_y)

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

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

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

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

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

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

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

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

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

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

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

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

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def calculate_adaptive_weight(
        self, nll_loss: torch.Tensor, g_loss: torch.Tensor, last_layer: torch.Tensor
    ) -> torch.Tensor:
        nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]

        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
        d_weight = d_weight * self.discriminator_weight
        return d_weight
//  [python]
def calculate_adaptive_weight(
        self, nll_loss: torch.Tensor, g_loss: torch.Tensor, last_layer: torch.Tensor
    ) -> torch.Tensor:
        nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]

        d_weight = torch.linalg.norm(nll_grads) / (torch.linalg.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
        d_weight = d_weight * self.discriminator_weight
        return d_weight
//  [python]
def forward(self, x):
      norm = torch.norm(x, p=2, dim=1, keepdim=True)
      clip = torch.clamp(norm, min=0)
      expand = clip.expand_as(x)
      return x / expand
//  [python]
def forward(self, x):
      norm = torch.linalg.norm(x, ord=2, dim=1, keepdim=True)
      clip = torch.clamp(norm, min=0)
      expand = clip.expand_as(x)
      return x / expand
//  [python]
def binary_search_gradient_estimation_variance(self, x):
        lower = self.attack_config["adaptive"]["bs_grad_var_lower"]
        upper = self.attack_config["adaptive"]["bs_grad_var_upper"]
        var = upper
        for _ in range(self.attack_config["adaptive"]["bs_grad_var_steps"]):
            mid = (lower + upper) / 2
            cache_hits = 0
            for _ in range(self.attack_config["adaptive"]["bs_grad_var_sample_size"]):
                noise = torch.randn_like(x).to(x.device)
                noise = noise / torch.norm(noise)
                noise = noise * mid
                noisy_img = x + noise
                noisy_img = torch.clamp(noisy_img, min=0, max=1)
                probs, is_cache = self.model(noisy_img)
                if is_cache[0]:
                    cache_hits += 1
            if cache_hits / self.attack_config["adaptive"]["bs_grad_var_sample_size"] \
                    <= self.attack_config["adaptive"]["bs_grad_var_hit_rate"]:
                var = mid
                upper = mid
            else:
                lower = mid
            print(
                f"Var : {var:.6f} | Cache Hits : {cache_hits}/{self.attack_config['adaptive']['bs_grad_var_sample_size']}")
        return var
//  [python]
def binary_search_gradient_estimation_variance(self, x):
        lower = self.attack_config["adaptive"]["bs_grad_var_lower"]
        upper = self.attack_config["adaptive"]["bs_grad_var_upper"]
        var = upper
        for _ in range(self.attack_config["adaptive"]["bs_grad_var_steps"]):
            mid = (lower + upper) / 2
            cache_hits = 0
            for _ in range(self.attack_config["adaptive"]["bs_grad_var_sample_size"]):
                noise = torch.randn_like(x).to(x.device)
                noise = noise / torch.linalg.norm(noise)
                noise = noise * mid
                noisy_img = x + noise
                noisy_img = torch.clamp(noisy_img, min=0, max=1)
                probs, is_cache = self.model(noisy_img)
                if is_cache[0]:
                    cache_hits += 1
            if cache_hits / self.attack_config["adaptive"]["bs_grad_var_sample_size"] \
                    <= self.attack_config["adaptive"]["bs_grad_var_hit_rate"]:
                var = mid
                upper = mid
            else:
                lower = mid
            print(
                f"Var : {var:.6f} | Cache Hits : {cache_hits}/{self.attack_config['adaptive']['bs_grad_var_sample_size']}")
        return var
//  [python]
def rotate_rodriguez(kp3d, rvec):
    '''
    kp3d: [b,vn, 3]
    rvec: [b,3]
    '''
    if len(kp3d.shape) == 2:
        kp3d = kp3d.unsqueeze(0)
    rvec = torch.atleast_2d(rvec)
    vn = kp3d.shape[1]
    theta = torch.norm(rvec)
    k = (rvec/theta).unsqueeze(2) # b,3,1
    k_T = k.permute(0,2,1) # b,1,3
    ctheta = torch.cos(theta)
    stheta = torch.sin(theta)
    rotated = kp3d * ctheta + torch.cross(k_T.expand(-1,vn,-1), kp3d, dim=2) * stheta + (k * torch.matmul(k_T, kp3d.permute(0,2,1))).permute(0,2,1) * (1-ctheta)
    return rotated
//  [python]
def rotate_rodriguez(kp3d, rvec):
    '''
    kp3d: [b,vn, 3]
    rvec: [b,3]
    '''
    if len(kp3d.shape) == 2:
        kp3d = kp3d.unsqueeze(0)
    rvec = torch.atleast_2d(rvec)
    vn = kp3d.shape[1]
    theta = torch.linalg.norm(rvec)
    k = (rvec/theta).unsqueeze(2) # b,3,1
    k_T = k.permute(0,2,1) # b,1,3
    ctheta = torch.cos(theta)
    stheta = torch.sin(theta)
    rotated = kp3d * ctheta + torch.cross(k_T.expand(-1,vn,-1), kp3d, dim=2) * stheta + (k * torch.matmul(k_T, kp3d.permute(0,2,1))).permute(0,2,1) * (1-ctheta)
    return rotated
//  [python]
def round_tt(
            self,
            eps: float = 1e-14,
            rmax: Optional[Union[int, Sequence[int]]] = None,
            algorithm: Optional[str] = 'svd',
            verbose: Optional[bool] = False):

        """
        Tries to recompress this tensor in place by reducing its TT ranks.

        Note: this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param eps: this relative error will not be exceeded
        :param rmax: all ranks should be rmax at most (default: no limit)
        :param algorithm: 'svd' (default) or 'eig'. The latter can be faster, but less accurate
        :param verbose:
        """

        N = self.dim()
        if not hasattr(rmax, '__len__'):
            rmax = [rmax] * (N - 1)
        assert len(rmax) == N - 1

        self._cp_to_tt()
        start = time.time()
        self.orthogonalize(N - 1)  # Make everything left-orthogonal
        if verbose:
            print('Orthogonalization time:', time.time() - start)
        if self.batch:
            delta = None
        else:
            delta = eps / max(1, torch.sqrt(torch.tensor([N - 1], dtype=torch.float64, device=self.cores[-1].device))) * torch.norm(self.cores[-1])
            delta = delta.item()

        for mu in range(N - 1, 0, -1):
            M = tn.right_unfolding(self.cores[mu], batch=self.batch)
            left, right = tn.truncated_svd(M, delta=delta, rmax=rmax[mu - 1], left_ortho=False, algorithm=algorithm, verbose=verbose, batch=self.batch)

            if self.batch:
                self.cores[mu] = right.reshape([self.cores[mu].shape[0], -1, self.cores[mu].shape[2], self.cores[mu].shape[3]])
                self.cores[mu - 1] = torch.einsum('bijk,bkl->bijl', (self.cores[mu - 1], left))  # Pass factor to the left
            else:
                self.cores[mu] = right.reshape([-1, self.cores[mu].shape[1], self.cores[mu].shape[2]])
                self.cores[mu - 1] = torch.einsum('ijk,kl', (self.cores[mu - 1], left))
//  [python]
def round_tt(
            self,
            eps: float = 1e-14,
            rmax: Optional[Union[int, Sequence[int]]] = None,
            algorithm: Optional[str] = 'svd',
            verbose: Optional[bool] = False):

        """
        Tries to recompress this tensor in place by reducing its TT ranks.

        Note: this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param eps: this relative error will not be exceeded
        :param rmax: all ranks should be rmax at most (default: no limit)
        :param algorithm: 'svd' (default) or 'eig'. The latter can be faster, but less accurate
        :param verbose:
        """

        N = self.dim()
        if not hasattr(rmax, '__len__'):
            rmax = [rmax] * (N - 1)
        assert len(rmax) == N - 1

        self._cp_to_tt()
        start = time.time()
        self.orthogonalize(N - 1)  # Make everything left-orthogonal
        if verbose:
            print('Orthogonalization time:', time.time() - start)
        if self.batch:
            delta = None
        else:
            delta = eps / max(1, torch.sqrt(torch.tensor([N - 1], dtype=torch.float64, device=self.cores[-1].device))) * torch.linalg.norm(self.cores[-1])
            delta = delta.item()

        for mu in range(N - 1, 0, -1):
            M = tn.right_unfolding(self.cores[mu], batch=self.batch)
            left, right = tn.truncated_svd(M, delta=delta, rmax=rmax[mu - 1], left_ortho=False, algorithm=algorithm, verbose=verbose, batch=self.batch)

            if self.batch:
                self.cores[mu] = right.reshape([self.cores[mu].shape[0], -1, self.cores[mu].shape[2], self.cores[mu].shape[3]])
                self.cores[mu - 1] = torch.einsum('bijk,bkl->bijl', (self.cores[mu - 1], left))  # Pass factor to the left
            else:
                self.cores[mu] = right.reshape([-1, self.cores[mu].shape[1], self.cores[mu].shape[2]])
                self.cores[mu - 1] = torch.einsum('ijk,kl', (self.cores[mu - 1], left))
//  [python]
def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor:
        """Spherical Linear intERPolation.

        Args:
            x0 (`torch.Tensor`):
                The first tensor to interpolate between.
            x1 (`torch.Tensor`):
                Second tensor to interpolate between.
            alpha (`float`):
                Interpolation between 0 and 1

        Returns:
            `torch.Tensor`:
                The interpolated tensor.
        """

        theta = acos(torch.dot(torch.flatten(x0), torch.flatten(x1)) / torch.norm(x0) / torch.norm(x1))
        return sin((1 - alpha) * theta) * x0 / sin(theta) + sin(alpha * theta) * x1 / sin(theta)
//  [python]
def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor:
        """Spherical Linear intERPolation.

        Args:
            x0 (`torch.Tensor`):
                The first tensor to interpolate between.
            x1 (`torch.Tensor`):
                Second tensor to interpolate between.
            alpha (`float`):
                Interpolation between 0 and 1

        Returns:
            `torch.Tensor`:
                The interpolated tensor.
        """

        theta = acos(torch.dot(torch.flatten(x0), torch.flatten(x1)) / torch.linalg.norm(x0) / torch.linalg.norm(x1))
        return sin((1 - alpha) * theta) * x0 / sin(theta) + sin(alpha * theta) * x1 / sin(theta)
//  [python]
def single_index(n,d,n1,sigma=0.1,nonlin=relu):
    X = unit_Gaussian(n,d)
    X_test = unit_Gaussian(n1,d)
    theta = torch.FloatTensor(d,1).normal_(mean=0,std=1).to(device)  
    theta = theta / torch.norm(theta)
    if sigma == 0:
        return X, nonlin(X@theta), X_test, nonlin(X_test@theta)
    else:
        return X, nonlin(X@theta) + torch.FloatTensor(n,1).normal_(mean=0,std=sigma).to(device), X_test, nonlin(X_test@theta)
//  [python]
def single_index(n,d,n1,sigma=0.1,nonlin=relu):
    X = unit_Gaussian(n,d)
    X_test = unit_Gaussian(n1,d)
    theta = torch.FloatTensor(d,1).normal_(mean=0,std=1).to(device)  
    theta = theta / torch.linalg.norm(theta)
    if sigma == 0:
        return X, nonlin(X@theta), X_test, nonlin(X_test@theta)
    else:
        return X, nonlin(X@theta) + torch.FloatTensor(n,1).normal_(mean=0,std=sigma).to(device), X_test, nonlin(X_test@theta)
//  [python]
def ineq_resid(self, X, Y):
        res = []
        """
         ||G^T y + h||_2 <= C^Ty+d
         G: m * k * n
         h: m * k
         y: m * n
         C: m * n
         d: m * 1
        """
        q = torch.norm(torch.matmul(self.G, Y.T).permute(2, 0, 1) + self.h.unsqueeze(0), dim=-1, p=2)
        p = torch.matmul(Y, self.C.T) + self.d
        res = q - p
        l = self.L - Y
        u = Y - self.U
        resids = torch.cat([res, l, u], dim=1)
        return torch.clamp(resids, 0)
//  [python]
def ineq_resid(self, X, Y):
        res = []
        """
         ||G^T y + h||_2 <= C^Ty+d
         G: m * k * n
         h: m * k
         y: m * n
         C: m * n
         d: m * 1
        """
        q = torch.linalg.norm(torch.matmul(self.G, Y.T).permute(2, 0, 1) + self.h.unsqueeze(0), dim=-1, ord=2)
        p = torch.matmul(Y, self.C.T) + self.d
        res = q - p
        l = self.L - Y
        u = Y - self.U
        resids = torch.cat([res, l, u], dim=1)
        return torch.clamp(resids, 0)
//  [python]
def resolve_equiv_natives(xs, natstack, maskstack):
    if (len(natstack.shape)==4):
        return natstack, maskstack
    if (natstack.shape[1]==1):
        return natstack[:,0,...], maskstack[:,0,...]
    dx = torch.norm( xs[:,None,:,None,1,:]-xs[:,None,None,:,1,:], dim=-1)
    dnat = torch.norm( natstack[:,:,:,None,1,:]-natstack[:,:,None,:,1,:], dim=-1)
    delta = torch.sum( torch.abs(dnat-dx), dim=(-2,-1))
    return natstack[:,torch.argmin(delta),...], maskstack[:,torch.argmin(delta),...]
//  [python]
def resolve_equiv_natives(xs, natstack, maskstack):
    if (len(natstack.shape)==4):
        return natstack, maskstack
    if (natstack.shape[1]==1):
        return natstack[:,0,...], maskstack[:,0,...]
    dx = torch.linalg.norm( xs[:,None,:,None,1,:]-xs[:,None,None,:,1,:], dim=-1)
    dnat = torch.linalg.norm( natstack[:,:,:,None,1,:]-natstack[:,:,None,:,1,:], dim=-1)
    delta = torch.sum( torch.abs(dnat-dx), dim=(-2,-1))
    return natstack[:,torch.argmin(delta),...], maskstack[:,torch.argmin(delta),...]
//  [python]
def multi_index(n,d,n1,sigma=0.1,ind=1,nonlin=relu):
    X = unit_Gaussian(n,d)
    X_test = unit_Gaussian(n1,d)
    theta = torch.FloatTensor(d,ind).normal_(mean=0,std=1).to(device) 
    theta = theta / torch.norm(theta)
    a = torch.ones([ind,1]).to(device) / (ind**0.5)
    Y = nonlin(X@theta)@a
    Y_test = nonlin(X_test@theta)@a
    if sigma == 0:
        return X, Y, X_test, Y_test
    else:
        return X, Y + torch.FloatTensor(n,1).normal_(mean=0,std=sigma).to(device), X_test, Y_test
//  [python]
def multi_index(n,d,n1,sigma=0.1,ind=1,nonlin=relu):
    X = unit_Gaussian(n,d)
    X_test = unit_Gaussian(n1,d)
    theta = torch.FloatTensor(d,ind).normal_(mean=0,std=1).to(device) 
    theta = theta / torch.linalg.norm(theta)
    a = torch.ones([ind,1]).to(device) / (ind**0.5)
    Y = nonlin(X@theta)@a
    Y_test = nonlin(X_test@theta)@a
    if sigma == 0:
        return X, Y, X_test, Y_test
    else:
        return X, Y + torch.FloatTensor(n,1).normal_(mean=0,std=sigma).to(device), X_test, Y_test
//  [python]
def ineq_resid(self, X, Y):
        res = []
        """
        Y>>0 -> xYx > 0
        """
        Ym = Y.view(Y.shape[0], self.ymdim, self.ymdim).permute(0, 2, 1)  # batch * n * n

        ### sample-based methods
        num_sample = 1024
        est = torch.randn(size=(1, self.ymdim, num_sample), device=X.device)  # batch * n * k
        est = est / torch.norm(est, dim=1, p=2, keepdim=True)
        pel = torch.matmul(Ym, est)  # batch * n * k
        pel = torch.multiply(pel, est).sum(1)  # batch * n * k
        res = -1*pel.view(-1, num_sample).sum(1, keepdim=True)

        l = self.L - Y
        u = Y - self.U
        resids = torch.cat([res, l, u], dim=1)
        return torch.clamp(resids, 0)
//  [python]
def ineq_resid(self, X, Y):
        res = []
        """
        Y>>0 -> xYx > 0
        """
        Ym = Y.view(Y.shape[0], self.ymdim, self.ymdim).permute(0, 2, 1)  # batch * n * n

        ### sample-based methods
        num_sample = 1024
        est = torch.randn(size=(1, self.ymdim, num_sample), device=X.device)  # batch * n * k
        est = est / torch.linalg.norm(est, dim=1, keepdim=True)
        pel = torch.matmul(Ym, est)  # batch * n * k
        pel = torch.multiply(pel, est).sum(1)  # batch * n * k
        res = -1*pel.view(-1, num_sample).sum(1, keepdim=True)

        l = self.L - Y
        u = Y - self.U
        resids = torch.cat([res, l, u], dim=1)
        return torch.clamp(resids, 0)
//  [python]
def dsacstar_loss(pred_X, gt_X, pixel_grid_crop, T, K, config):
    """
    :param pred_X: B,3,H,W
    :param gt_X: B,N,3
    :param pixel_grid_crop: N,2
    :param T: B,3,4
    :param K: B,3,3
    :return:
    """
    R = T[:, :3, :3]
    t = T[:, :3, 3].squeeze(-1)

    pred_X = pred_X.view(T.size(0), 3, -1).permute(0,2,1)
    B, N, _ = pred_X.size()
    X_camera = batched_transpose(R, t, pred_X)   
    projected_2d, _ = batched_pi(K, X_camera)

    reprojection_error = projected_2d - pixel_grid_crop
    reprojection_error = reprojection_error.norm(2, 2)

    # check constraints
    invalid_min_depth = X_camera[:, :, 2] < config['mindepth']  
    invalid_repro = reprojection_error > config['hardclamp']  

    if gt_X is not None:
        gt_coords_mask = torch.abs(gt_X).sum(2) == 0

        target_camera = batched_transpose(R, t, gt_X)
        gt_coord_dist = torch.norm(X_camera - target_camera, dim=2, p=2)

        invalid_gt_distance = gt_coord_dist > config['itol'] 
        invalid_gt_distance[gt_coords_mask] = 0  

        valid_scene_coordinates = (invalid_min_depth + invalid_gt_distance + invalid_repro) == 0
    else:
        invalid_max_depth = X_camera[:, :, 2] > config['maxdepth']
        valid_scene_coordinates = (invalid_min_depth + invalid_max_depth + invalid_repro) == 0

    num_valid_sc = int(valid_scene_coordinates.sum())

    loss = 0

    if num_valid_sc > 0:
        
        reprojection_error = reprojection_error[valid_scene_coordinates]

        loss_l1 = reprojection_error[reprojection_error <= config['softclamp']]
        loss_sqrt = reprojection_error[reprojection_error > config['softclamp']]
        loss_sqrt = torch.sqrt(config['softclamp'] * loss_sqrt)

        loss += (loss_l1.sum() + loss_sqrt.sum())

    if num_valid_sc < B*N:

        invalid_scene_coordinates = (valid_scene_coordinates == 0)

        if gt_X is not None:
            invalid_scene_coordinates[gt_coords_mask] = 0
            loss += gt_coord_dist[invalid_scene_coordinates].sum()
        else:
            depth = config['targetdepth'] * torch.ones(B,N,1).to(device)
            pixel_grid_crop_batched = pixel_grid_crop.expand(B,-1,-1)
            target_camera = batched_pi_inv(K, pixel_grid_crop_batched, depth)

            loss += torch.abs(
                X_camera[invalid_scene_coordinates] - target_camera[invalid_scene_coordinates]).sum()

    loss /= B*N
    num_valid_sc /= B*N

    return loss, num_valid_sc
//  [python]
def dsacstar_loss(pred_X, gt_X, pixel_grid_crop, T, K, config):
    """
    :param pred_X: B,3,H,W
    :param gt_X: B,N,3
    :param pixel_grid_crop: N,2
    :param T: B,3,4
    :param K: B,3,3
    :return:
    """
    R = T[:, :3, :3]
    t = T[:, :3, 3].squeeze(-1)

    pred_X = pred_X.view(T.size(0), 3, -1).permute(0,2,1)
    B, N, _ = pred_X.size()
    X_camera = batched_transpose(R, t, pred_X)   
    projected_2d, _ = batched_pi(K, X_camera)

    reprojection_error = projected_2d - pixel_grid_crop
    reprojection_error = torch.linalg.norm(reprojection_error, 2, 2)

    # check constraints
    invalid_min_depth = X_camera[:, :, 2] < config['mindepth']  
    invalid_repro = reprojection_error > config['hardclamp']  

    if gt_X is not None:
        gt_coords_mask = torch.abs(gt_X).sum(2) == 0

        target_camera = batched_transpose(R, t, gt_X)
        gt_coord_dist = torch.linalg.norm(X_camera - target_camera, dim=2, ord=2)

        invalid_gt_distance = gt_coord_dist > config['itol'] 
        invalid_gt_distance[gt_coords_mask] = 0  

        valid_scene_coordinates = (invalid_min_depth + invalid_gt_distance + invalid_repro) == 0
    else:
        invalid_max_depth = X_camera[:, :, 2] > config['maxdepth']
        valid_scene_coordinates = (invalid_min_depth + invalid_max_depth + invalid_repro) == 0

    num_valid_sc = int(valid_scene_coordinates.sum())

    loss = 0

    if num_valid_sc > 0:
        
        reprojection_error = reprojection_error[valid_scene_coordinates]

        loss_l1 = reprojection_error[reprojection_error <= config['softclamp']]
        loss_sqrt = reprojection_error[reprojection_error > config['softclamp']]
        loss_sqrt = torch.sqrt(config['softclamp'] * loss_sqrt)

        loss += (loss_l1.sum() + loss_sqrt.sum())

    if num_valid_sc < B*N:

        invalid_scene_coordinates = (valid_scene_coordinates == 0)

        if gt_X is not None:
            invalid_scene_coordinates[gt_coords_mask] = 0
            loss += gt_coord_dist[invalid_scene_coordinates].sum()
        else:
            depth = config['targetdepth'] * torch.ones(B,N,1).to(device)
            pixel_grid_crop_batched = pixel_grid_crop.expand(B,-1,-1)
            target_camera = batched_pi_inv(K, pixel_grid_crop_batched, depth)

            loss += torch.abs(
                X_camera[invalid_scene_coordinates] - target_camera[invalid_scene_coordinates]).sum()

    loss /= B*N
    num_valid_sc /= B*N

    return loss, num_valid_sc
//  [python]
def forecasting_loss_and_metrics(gt: Tensor, pred: Tensor, prefix='pred') -> (Tensor, dict[str, Tensor]):
    # Compute state squared error over time and the infinite norm of the state dimension over time.
    l2_loss = torch.norm(gt - pred, p=2, dim=-1)
    metrics = {}
    metrics[f'{prefix}_loss_t'] = l2_loss
    return l2_loss.mean(), metrics
//  [python]
def forecasting_loss_and_metrics(gt: Tensor, pred: Tensor, prefix='pred') -> (Tensor, dict[str, Tensor]):
    # Compute state squared error over time and the infinite norm of the state dimension over time.
    l2_loss = torch.linalg.norm(gt - pred, ord=2, dim=-1)
    metrics = {}
    metrics[f'{prefix}_loss_t'] = l2_loss
    return l2_loss.mean(), metrics
//  [python]
def transformation_error(pred_trans, gt_trans):
    pred_R = pred_trans[:3, :3]
    gt_R = gt_trans[:3, :3]
    pred_t = pred_trans[:3, 3:4]
    gt_t = gt_trans[:3, 3:4]
    tr = torch.trace(pred_R.T @ gt_R)
    RE = torch.acos(torch.clamp((tr - 1) / 2.0, min=-1, max=1)) * 180 / np.pi
    TE = torch.norm(pred_t - gt_t) * 100
    return RE, TE
//  [python]
def transformation_error(pred_trans, gt_trans):
    pred_R = pred_trans[:3, :3]
    gt_R = gt_trans[:3, :3]
    pred_t = pred_trans[:3, 3:4]
    gt_t = gt_trans[:3, 3:4]
    tr = torch.trace(pred_R.T @ gt_R)
    RE = torch.acos(torch.clamp((tr - 1) / 2.0, min=-1, max=1)) * 180 / np.pi
    TE = torch.linalg.norm(pred_t - gt_t) * 100
    return RE, TE
//  [python]
def euclidean_dist(eigen_values):
    roll = opt.batchSize#eigen_values.size()[0]
    total_dist = 0
    for i in range(roll):
        for j in range(roll):
            total_dist += torch.norm(eigen_values[i, :] - eigen_values[j, :], float("inf"))
#    print('Euclidean: ', 3*torch.exp(-1/total_dist))
#    return 3*torch.exp(-1/total_dist)
    return 100000000/total_dist
//  [python]
def euclidean_dist(eigen_values):
    roll = opt.batchSize#eigen_values.size()[0]
    total_dist = 0
    for i in range(roll):
        for j in range(roll):
            total_dist += torch.linalg.norm(eigen_values[i, :] - eigen_values[j, :], float("inf"))
#    print('Euclidean: ', 3*torch.exp(-1/total_dist))
#    return 3*torch.exp(-1/total_dist)
    return 100000000/total_dist
//  [python]
def post_refinement(initial_trans, src_kpts, tgt_kpts, iters, inlier_threshold = 0.1, weights=None):
    pre_inlier_count = 0
    for i in range(iters):
        pred_tgt = transform(src_kpts, initial_trans)
        L2_dis = torch.norm(pred_tgt - tgt_kpts, dim=-1)
        pred_inlier = (L2_dis < inlier_threshold)[0]
        inlier_count = torch.sum(pred_inlier)
        if inlier_count <= pre_inlier_count:
            break
        pre_inlier_count = inlier_count
        initial_trans = rigid_transform_3d(
            A=src_kpts[:, pred_inlier, :],
            B=tgt_kpts[:, pred_inlier, :],
            weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier]
        )
    return initial_trans
//  [python]
def post_refinement(initial_trans, src_kpts, tgt_kpts, iters, inlier_threshold = 0.1, weights=None):
    pre_inlier_count = 0
    for i in range(iters):
        pred_tgt = transform(src_kpts, initial_trans)
        L2_dis = torch.linalg.norm(pred_tgt - tgt_kpts, dim=-1)
        pred_inlier = (L2_dis < inlier_threshold)[0]
        inlier_count = torch.sum(pred_inlier)
        if inlier_count <= pre_inlier_count:
            break
        pre_inlier_count = inlier_count
        initial_trans = rigid_transform_3d(
            A=src_kpts[:, pred_inlier, :],
            B=tgt_kpts[:, pred_inlier, :],
            weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier]
        )
    return initial_trans
//  [python]
def __init__(self, pca_dict, sparsity):
        normed_dict = pca_dict / torch.norm(pca_dict, dim=-1)[:, None]
        self.pca_dict = normed_dict
        self.sparsity = sparsity
        self.n_feats, self.activation_size = self.pca_dict.shape
//  [python]
def __init__(self, pca_dict, sparsity):
        normed_dict = pca_dict / torch.linalg.norm(pca_dict, dim=-1)[:, None]
        self.pca_dict = normed_dict
        self.sparsity = sparsity
        self.n_feats, self.activation_size = self.pca_dict.shape
//  [python]
def _compute_norm(t, n, dim):
    r"""Compute the L_n-norm of a tensor along all dimensions except for the specified dimension.

    The L_n-norm will be computed across all entries in tensor `t` along all dimension
    except for the one identified by dim.
    Example: if `t` is of shape, say, 3x2x4 and dim=2 (the last dim),
    then norm will have Size [4], and each entry will represent the
    `L_n`-norm computed using the 3x2=6 entries for each of the 4 channels.

    Args:
        t (torch.Tensor): tensor representing the parameter to prune
        n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid
            entries for argument p in torch.norm
        dim (int): dim identifying the channels to prune

    Returns:
        norm (torch.Tensor): L_n norm computed across all dimensions except
            for `dim`. By construction, `norm.shape = t.shape[-1]`.
    """
    # dims = all axes, except for the one identified by `dim`
    dims = list(range(t.dim()))
    # convert negative indexing
    if dim < 0:
        dim = dims[dim]
    dims.remove(dim)

    norm = torch.norm(t, p=n, dim=dims)
    return norm