jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def empirical_lin_inverse_projector(self, state: Tensor, obs_state: Tensor):
        """ Compute the empirical inverse projector from the observable state to the pre-processed state.

        Args:
            state: (batch, state_dim) Tensor containing the pre-processed state.
            obs_state: (batch, obs_state_dim) Tensor containing the observable state.
        Returns:
            A: (state_dim, obs_state_dim) Tensor containing the empirical inverse projector.
            B: (state_dim, 1) Tensor containing the empirical inverse projector bias.
            rec_error: Scalar tensor containing the reconstruction error or "residual".
        """
        # Inverse projector is computed from the observable state to the pre-processed state
        pre_state = self.pre_process_state(state)
        X = obs_state.T  # (obs_state_dim, n_samples)
        Y = pre_state.T  # (state_dim, n_samples)
        A, B = full_rank_lstsq(X, Y, bias=True)

        rec_error = torch.nn.functional.mse_loss(A @ X + B, Y)

        metrics = dict(inverse_projector_rank=torch.linalg.matrix_rank(A.detach()).to(torch.float),
                       inverse_projector_cond_num=torch.linalg.cond(A.detach()).to(torch.float),
                       inverse_projector_error=rec_error)

        return A, B, metrics
//  [python]
def z2rank(mat):
    '''Calculate Z2 rank of a binary matrix.
    Parameters:
    mat: int matrix - input binary matrix.
        caller must ensure that mat contains only 0 and 1.
        mat is destroyed upon output!
    Returns:
    r: int - rank of the matrix under Z2 algebra.'''
    return torch.matrix_rank(mat.to(torch.float32))
//  [python]
def z2rank(mat):
    '''Calculate Z2 rank of a binary matrix.
    Parameters:
    mat: int matrix - input binary matrix.
        caller must ensure that mat contains only 0 and 1.
        mat is destroyed upon output!
    Returns:
    r: int - rank of the matrix under Z2 algebra.'''
    return torch.linalg.matrix_rank(mat.to(torch.float32))
//  [python]
def test_matrix_rank_out_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        out = torch.empty(0, dtype=torch.bool, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got result with dtype Bool"):
            torch.linalg.matrix_rank(a, out=out)

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.matrix_rank(a, out=out)

        # if out tensor with wrong shape is passed a warning is given
        with warnings.catch_warnings(record=True) as w:
            out = torch.empty(3, dtype=dtype, device=device)
            # Trigger warning
            torch.linalg.matrix_rank(a, out=out)
            # Check warning occurs
            self.assertEqual(len(w), 1)
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))
//  [python]
def test_matrix_rank_tol(self, device, dtype):

        def run_test_tol(shape0, shape1, batch):
            a = make_tensor((*batch, shape0, shape1), dtype=dtype, device=device)
            # Check against NumPy output
            # Test float tol, and specific value for each matrix
            tolerances = [float(torch.rand(1)), ]
            # Test different types of tol tensor
            for tol_type in all_types():
                tolerances.append(make_tensor(a.shape[:-2], dtype=tol_type, device=device, low=0))
            # Test broadcasting of tol
            if a.ndim > 2:
                tolerances.append(make_tensor(a.shape[-3], dtype=torch.float32, device=device, low=0))
            for tol in tolerances:
                actual = torch.linalg.matrix_rank(a, tol=tol)
                numpy_tol = tol if isinstance(tol, float) else tol.cpu().numpy()
                expected = np.linalg.matrix_rank(a.cpu().numpy(), tol=numpy_tol)
                self.assertEqual(actual, expected)

        shapes = (3, 13)
        batches = ((), (0, ), (4, ), (3, 5, ))
        for (shape0, shape1), batch in zip(itertools.product(shapes, reversed(shapes)), batches):
            run_test_tol(shape0, shape1, batch)
//  [python]
def lstsq(self, A, Y, lamb=0.0):
        """
        Differentiable least square
        :param A: m x n
        :param Y: n x 1
        """
        cols = A.shape[1]
        if np.isinf(A.data.cpu().numpy()).any():
            raise RuntimeException("Infinity in least squares")

        # Assuming A to be full column rank
        if cols == torch.linalg.matrix_rank(A):
            # Full column rank
            q, r = torch.linalg.qr(A)
            x = torch.inverse(r) @ q.transpose(1, 0) @ Y
        else:
            # rank(A) < n, do regularized least square.
            AtA = A.transpose(1, 0) @ A

            # get the smallest lambda that suits our purpose, so that error in
            # results minimized.
            with torch.no_grad():
                lamb = best_lambda(AtA)
            A_dash = AtA + lamb * torch.eye(cols, device=A.get_device())
            Y_dash = A.transpose(1, 0) @ Y

            # if it still doesn't work, just set the lamb to be very high value.
            x = self.lstsq(A_dash, Y_dash, 1)
        return x
//  [python]
def _length_impl(self, pt_a: Point, pt_b: Point) -> Scalar:
        # If rank-deficient, return infinity and be done early
        if torch.linalg.matrix_rank(pt_a) < self.shape[0] or torch.linalg.matrix_rank(pt_b) < self.shape[0]:
            return torch.tensor([float('inf')], dtype=pt_a.dtype, device=pt_a.device)
        # TODO - do we need eps and rank checks if we use a pseudo-inverse instead? Or will eigs be zero
        # and therefore dist --> infinity?
        inv_a_half = inv_matrix_sqrt(pt_a)
        x_inv_y = inv_a_half @ pt_b @ inv_a_half
        log_eigs = torch.log(torch.linalg.eigvals(x_inv_y).real)
        return torch.sqrt(torch.sum(log_eigs**2))
//  [python]
def create_full_rank_data() -> Tensor:
    # Create a full-rank data matrix of size (N, data_dimension).
    # Since torch.rand() is very likely to return a full-rank matrix,
    # we use it to achieve this (after setting seed).
    X = torch.rand(N, data_dimension)
    assert torch.matrix_rank(X) == data_dimension
    return X
//  [python]
def create_full_rank_data() -> Tensor:
    # Create a full-rank data matrix of size (N, data_dimension).
    # Since torch.rand() is very likely to return a full-rank matrix,
    # we use it to achieve this (after setting seed).
    X = torch.rand(N, data_dimension)
    assert torch.linalg.matrix_rank(X) == data_dimension
    return X
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.opts.lambda_reg * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I
        orig_rank = torch.matrix_rank(HTH).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self.device)
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.opts.lambda_reg * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I
        # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self.device)
        return W, tag
//  [python]
def main():
    torch.set_printoptions(precision=3)
    ckpt_dir = "./artifacts/LLAMA/llama-7b-hf"

    generator = hf_load(ckpt_dir)
    # param_dict = {}
    for name, layer in generator.model.named_parameters():
        param = layer.data.detach()
        rank = -1
        if param.ndim == 2:
            rank = torch.matrix_rank(param)
            # svd_vals = torch.linalg.svdvals(param)
            # param_dict[name] = svd_vals
        logger.debug(f"Layer: {name} has shape {layer.shape} with rank {rank} with dtype: {layer.dtype}")
//  [python]
def main():
    torch.set_printoptions(precision=3)
    ckpt_dir = "./artifacts/LLAMA/llama-7b-hf"

    generator = hf_load(ckpt_dir)
    # param_dict = {}
    for name, layer in generator.model.named_parameters():
        param = layer.data.detach()
        rank = -1
        if param.ndim == 2:
            rank = torch.linalg.matrix_rank(param)
            # svd_vals = torch.linalg.svdvals(param)
            # param_dict[name] = svd_vals
        logger.debug(f"Layer: {name} has shape {layer.shape} with rank {rank} with dtype: {layer.dtype}")
//  [python]
def checkMatrixRank(X, name):
    rank = torch.matrix_rank(X)
    print(f"{name} matrix rank: {rank}")
    if rank != X.shape[0]:
        # raise ValueError(f"{name} matrix is not full rank ({rank} / {X.shape[0]})")
        pass
//  [python]
def checkMatrixRank(X, name):
    rank = torch.linalg.matrix_rank(X)
    print(f"{name} matrix rank: {rank}")
    if rank != X.shape[0]:
        # raise ValueError(f"{name} matrix is not full rank ({rank} / {X.shape[0]})")
        pass
//  [python]
def create_low_rank_data() -> Tensor:
    # Create a low-rank data matrix of size (N, data_dimension) and rank = low_rank.
    A = torch.rand(small_rank - 1, data_dimension)
    assert torch.linalg.matrix_rank(A) == small_rank - 1
    B = torch.ones(N - small_rank + 1, data_dimension)
    X = torch.concat((A, B), dim=0)
    assert torch.linalg.matrix_rank(X) == small_rank
    return X
//  [python]
def build_graph(self, config):
        x = self.variable(name='x', shape=config.x_shape, dtype=config.x_dtype)
        result = torch.matrix_rank(
            input=x, tol=config.tol, hermitian=config.hermitian)

        self.feed_list = [x]
        self.fetch_list = [result]
//  [python]
def build_graph(self, config):
        x = self.variable(name='x', shape=config.x_shape, dtype=config.x_dtype)
        result = torch.linalg.matrix_rank(
            input=x, tol=config.tol, hermitian=config.hermitian)

        self.feed_list = [x]
        self.fetch_list = [result]
//  [python]
def compute_hist1(self,idx=0):
        dens, bins = torch.histogram(self.IP[idx, :, :].flatten().cpu(), 250, density=True);
        T = self.IP[idx, :, :]; 
        T1 = (torch.real(T)>0.9999).type(torch.FloatTensor);
        nr_clust = torch.linalg.matrix_rank(T1, hermitian=False);
        T1L, T1V = torch.linalg.eig(T1);
        T1L = torch.abs(T1L)
        cluster_sizes = (T1L[T1L>0.001]);
        return dens, bins, cluster_sizes;
//  [python]
def count_clusters(self):
        batch, ntokens = self.args.batch, self.args.ntokens;
        
        list_clust = [];
        for i in range(batch):
            T = self.IP[i,:,:]; 
            T1 = (torch.real(T)>0.9999).type(torch.FloatTensor);
            nr_clust = torch.linalg.matrix_rank(T1, hermitian=False);
            #T1L, T1V = torch.linalg.eig(T1); 
            #cluster_sizes = (T1L[torch.abs(T1L)>1e-3])
            list_clust.append(nr_clust);
            
        return np.asarray(list_clust);
//  [python]
def __init__(self,
                 state_dim: Optional[int] = None,
                 state_rep: Optional[Representation] = None,
                 dt: Optional[Union[float, int]] = 1,
                 state_change_of_basis: Optional[Tensor] = None,
                 state_inv_change_of_basis: Optional[Tensor] = None):
        super().__init__()
        assert state_dim is not None or state_rep is not None, \
            "Either state_dim or state_rep must be provided"
        self.state_rep: Representation = state_rep
        self.state_dim = state_dim if state_dim is not None else state_rep.size
        self.dt = dt

        # Configure optional change of basis of the state.
        if state_change_of_basis is not None:
            assert state_change_of_basis.shape == (self.state_dim, self.state_dim), \
                f"Expected change of basis ({self.state_dim}, {self.state_dim}), got {state_change_of_basis.shape}"
            assert torch.linalg.matrix_rank(state_change_of_basis) == self.state_dim, \
                f"Defective state_change_of_basis with rank {torch.linalg.matrix_rank(state_change_of_basis)}"
            if state_inv_change_of_basis is None:
                state_inv_change_of_basis = torch.linalg.inv(state_change_of_basis)
            else:
                assert state_inv_change_of_basis.shape == state_change_of_basis.shape

        self.state_change_of_basis = state_change_of_basis
        self.state_inv_change_of_basis = state_inv_change_of_basis

        # Used to reshape the state trajectory from/back (batch * time, state_dim) to (batch, time, state_dim)
        self._batch_size = None
        log.info(f"Instantiating {self.__class__.__name__} with state_dim={self.state_dim} and dt={self.dt}")
//  [python]
def test_matrix_rank_atol(self, device, dtype):

        def run_test_atol(shape0, shape1, batch):
            a = make_tensor((*batch, shape0, shape1), dtype=dtype, device=device)
            # Check against NumPy output
            # Test float tol, and specific value for each matrix
            tolerances = [float(torch.rand(1)), ]
            # Test different types of tol tensor
            for tol_type in all_types():
                tolerances.append(make_tensor(a.shape[:-2], dtype=tol_type, device=device, low=0))
            # Test broadcasting of tol
            if a.ndim > 2:
                tolerances.append(make_tensor(a.shape[-3], dtype=torch.float32, device=device, low=0))
            for tol in tolerances:
                actual = torch.linalg.matrix_rank(a, atol=tol)
                actual_tol = torch.linalg.matrix_rank(a, tol=tol)
                self.assertEqual(actual, actual_tol)
                numpy_tol = tol if isinstance(tol, float) else tol.cpu().numpy()
                expected = np.linalg.matrix_rank(a.cpu().numpy(), tol=numpy_tol)
                self.assertEqual(actual, expected)

        shapes = (3, 13)
        batches = ((), (0, ), (4, ), (3, 5, ))
        for (shape0, shape1), batch in zip(itertools.product(shapes, reversed(shapes)), batches):
            run_test_atol(shape0, shape1, batch)
//  [python]
def test_matrix_rank_atol_rtol(self, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_arg = partial(make_fullrank, device=device, dtype=dtype)

        # creates a matrix with singular values rank=n and singular values in range [2/3, 3/2]
        # the singular values are 1 + 1/2, 1 - 1/3, 1 + 1/4, 1 - 1/5, ...
        n = 9
        a = make_arg(n, n)

        # test float and tensor variants
        for tol_value in [0.81, torch.tensor(0.81, device=device)]:
            # using rtol (relative tolerance) takes into account the largest singular value (1.5 in this case)
            result = torch.matrix_rank(a, rtol=tol_value)
            self.assertEqual(result, 2)  # there are 2 singular values above 1.5*0.81 = 1.215

            # atol is used directly to compare with singular values
            result = torch.matrix_rank(a, atol=tol_value)
            self.assertEqual(result, 7)  # there are 7 singular values above 0.81

            # when both are specified the maximum tolerance is used
            result = torch.matrix_rank(a, atol=tol_value, rtol=tol_value)
            self.assertEqual(result, 2)
//  [python]
def test_matrix_rank_atol_rtol(self, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_arg = partial(make_fullrank, device=device, dtype=dtype)

        # creates a matrix with singular values rank=n and singular values in range [2/3, 3/2]
        # the singular values are 1 + 1/2, 1 - 1/3, 1 + 1/4, 1 - 1/5, ...
        n = 9
        a = make_arg(n, n)

        # test float and tensor variants
        for tol_value in [0.81, torch.tensor(0.81, device=device)]:
            # using rtol (relative tolerance) takes into account the largest singular value (1.5 in this case)
            result = torch.linalg.matrix_rank(a, rtol=tol_value)
            self.assertEqual(result, 2)  # there are 2 singular values above 1.5*0.81 = 1.215

            # atol is used directly to compare with singular values
            result = torch.linalg.matrix_rank(a, atol=tol_value)
            self.assertEqual(result, 7)  # there are 7 singular values above 0.81

            # when both are specified the maximum tolerance is used
            result = torch.linalg.matrix_rank(a, atol=tol_value, rtol=tol_value)
            self.assertEqual(result, 2)
//  [python]
def matrix_rank(a: ArrayLike, tol=None, hermitian=False):
    a = _atleast_float_1(a)

    if a.ndim < 2:
        return int((a != 0).any())

    if tol is None:
        # follow https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1885
        atol = 0
        rtol = max(a.shape[-2:]) * torch.finfo(a.dtype).eps
    else:
        atol, rtol = tol, 0
    return torch.linalg.matrix_rank(a, atol=atol, rtol=rtol, hermitian=hermitian)
//  [python]
def get_inv_cov(x_dat, main_diag_only, cuda=False):
    """Get inverse of covariance matrix."""
    cols_x = x_dat.shape[1]
    if cuda:
        x_dat_t = torch.from_numpy(x_dat.T)
        x_dat_t = x_dat_t.to('cuda')
        try:
            cov_x = torch.cov(x_dat_t)
        except AttributeError:
            x_dat_t = x_dat_t.to(torch.float32)
            cov_x = torch.cov(x_dat_t)
        if main_diag_only:
            cov_x = cov_x * torch.eye(cols_x, device='cuda')   # only main diag
        rank_not_ok = True
        counter = 0
        while rank_not_ok:
            if counter == 20:
                cov_x = cov_x * torch.eye(cols_x, device='cuda')
            if counter > 20:
                cov_x_inv = torch.eye(cols_x, device='cuda')
                break
            if torch.matrix_rank(cov_x) < cols_x:
                cov_x += 0.5 * torch.diag(cov_x) * torch.eye(cols_x,
                                                             device='cuda')
                counter += 1
            else:
                cov_x_inv_t = torch.inverse(cov_x)
                rank_not_ok = False
        cov_x_inv_t = cov_x_inv_t.to('cpu')
        cov_x_inv = cov_x_inv_t.numpy()
    else:
        try:
            cov_x = np.cov(x_dat, rowvar=False)
        except AttributeError:
            cov_x = np.cov(x_dat.astype(np.float32), rowvar=False)
        if main_diag_only:
            cov_x = cov_x * np.eye(cols_x)   # only main diag
        rank_not_ok = True
        counter = 0
        while rank_not_ok:
            if counter == 20:
                cov_x *= np.eye(cols_x)
            if counter > 20:
                cov_x_inv = np.eye(cols_x)
                break
            if np.linalg.matrix_rank(cov_x) < cols_x:
                cov_x += 0.5 * np.diag(cov_x) * np.eye(cols_x)
                counter += 1
            else:
                cov_x_inv = np.linalg.inv(cov_x)
                rank_not_ok = False
    return cov_x_inv
//  [python]
def get_inv_cov(x_dat, main_diag_only, cuda=False):
    """Get inverse of covariance matrix."""
    cols_x = x_dat.shape[1]
    if cuda:
        x_dat_t = torch.from_numpy(x_dat.T)
        x_dat_t = x_dat_t.to('cuda')
        try:
            cov_x = torch.cov(x_dat_t)
        except AttributeError:
            x_dat_t = x_dat_t.to(torch.float32)
            cov_x = torch.cov(x_dat_t)
        if main_diag_only:
            cov_x = cov_x * torch.eye(cols_x, device='cuda')   # only main diag
        rank_not_ok = True
        counter = 0
        while rank_not_ok:
            if counter == 20:
                cov_x = cov_x * torch.eye(cols_x, device='cuda')
            if counter > 20:
                cov_x_inv = torch.eye(cols_x, device='cuda')
                break
            if torch.linalg.matrix_rank(cov_x) < cols_x:
                cov_x += 0.5 * torch.diag(cov_x) * torch.eye(cols_x,
                                                             device='cuda')
                counter += 1
            else:
                cov_x_inv_t = torch.linalg.inv(cov_x)
                rank_not_ok = False
        cov_x_inv_t = cov_x_inv_t.to('cpu')
        cov_x_inv = cov_x_inv_t.numpy()
    else:
        try:
            cov_x = np.cov(x_dat, rowvar=False)
        except AttributeError:
            cov_x = np.cov(x_dat.astype(np.float32), rowvar=False)
        if main_diag_only:
            cov_x = cov_x * np.eye(cols_x)   # only main diag
        rank_not_ok = True
        counter = 0
        while rank_not_ok:
            if counter == 20:
                cov_x *= np.eye(cols_x)
            if counter > 20:
                cov_x_inv = np.eye(cols_x)
                break
            if np.linalg.matrix_rank(cov_x) < cols_x:
                cov_x += 0.5 * np.diag(cov_x) * np.eye(cols_x)
                counter += 1
            else:
                cov_x_inv = np.linalg.inv(cov_x)
                rank_not_ok = False
    return cov_x_inv
//  [python]
def check_stats(self):

        self.update_IP();

        batch = self.args.batch;

        if self.args.plotdim:
            try:
                eigv = torch.linalg.eigvalsh(self.IP); # batch x ntoken
            except Exception as inst:
                print('[ERROR] linalg.eigvalsh failed: ', inst);
                dens = 0;
            else:
                # print('[debug] eigv = ', eigv);
                dens = (eigv > 0.001).flatten().sum() / batch;
        else:

            # Earlier versions had a bug where we also counted i=j particles into the sum.
            mask = ~torch.eye(self.args.ntokens).bool().expand(batch, -1, -1);
            inner_prods = self.IP[mask]; # this should be automatically flattened
            dens = (inner_prods > 0.999).sum() / len(inner_prods);


        ### This is a version that counts i=j particles as well. Use it to check for numerical instabilities
        #inner_prods = self.IP.flatten();
        #dens = (inner_prods > 0.999).sum() / len(inner_prods);
        
        if self.args.cluster_sizes:
            time0 = time.time();
            list_clust = [];

            for i in range(batch):
                T = self.IP[i,:,:]; 
                T1 = (torch.real(T)>0.9999).type(torch.FloatTensor);
                U, S, V = torch.linalg.svd(T1)
                nr_clust = torch.sum(S > 1e-5)
                #T1L, T1V = torch.linalg.eig(T1); 
                #cluster_sizes = (T1L[torch.abs(T1L)>1e-3])
                list_clust.append(nr_clust);

            median_clust = np.median(np.asarray(list_clust));
            time1 = time.time();

            print(f'[debug] check_stats() took {(time1-time0):.2f} sec');
        else:
            median_clust = 0;

        return dens, median_clust;
//  [python]
def check_stats(self):

        self.update_IP();

        batch = self.args.batch;

        if self.args.plotdim:
            try:
                eigv = torch.linalg.eigvalsh(self.IP); # batch x ntoken
            except Exception as inst:
                print('[ERROR] linalg.eigvalsh failed: ', inst);
                dens = 0;
            else:
                # print('[debug] eigv = ', eigv);
                dens = (eigv > 0.001).flatten().sum() / batch;
        else:

            # Earlier versions had a bug where we also counted i=j particles into the sum.
            mask = ~torch.eye(self.args.ntokens).bool().expand(batch, -1, -1);
            inner_prods = self.IP[mask]; # this should be automatically flattened
            dens = (inner_prods > 0.999).sum() / len(inner_prods);


        ### This is a version that counts i=j particles as well. Use it to check for numerical instabilities
        #inner_prods = self.IP.flatten();
        #dens = (inner_prods > 0.999).sum() / len(inner_prods);
        
        if self.args.cluster_sizes:
            time0 = time.time();
            list_clust = [];

            for i in range(batch):
                T = self.IP[i,:,:]; 
                T1 = (torch.real(T)>0.9999).type(torch.FloatTensor);
                nr_clust = torch.linalg.matrix_rank(T1, hermitian=False);
                #T1L, T1V = torch.linalg.eig(T1); 
                #cluster_sizes = (T1L[torch.abs(T1L)>1e-3])
                list_clust.append(nr_clust);

            median_clust = np.median(np.asarray(list_clust));
            time1 = time.time();

            print(f'[debug] check_stats() took {(time1-time0):.2f} sec');
        else:
            median_clust = 0;

        return dens, median_clust;
//  [python]
def lstq(self, A, Y, lamb=0.0, attempt=0):
        """
        Differentiable least square
        :param A: m x n
        :param Y: n x 1
        """
        # Assuming A to be full column rank
        cols = A.shape[1]
        try:
            rank = torch.linalg.matrix_rank(A)
        except:
            rank = -1
        if cols == rank:
            q, r = torch.linalg.qr(A)
            x = torch.inverse(r) @ q.T @ Y
        else:
            if attempt > self.max_attempts:
                print(f"Can't compute pseudoinverse, attempt: {attempt}; lambda: {lamb}")
                raise Exception("Can't compute pseudoinverse")
            A_dash = A.permute(1, 0) @ A + lamb * torch.eye(cols)
            Y_dash = A.permute(1, 0) @ Y
            x = self.lstq(A_dash, Y_dash, lamb=lamb, attempt=attempt + 1)
        return x
//  [python]
def get_plane_parameter(pc, pc_w):
    # min least square
    pc_w = pc_w / torch.max(pc_w)
    # min least square
    n = pc.shape[0]
    A = torch.cat([pc[:, :2], torch.ones([n, 1], device=pc.device)], dim=-1)
    b = pc[:, 2].view(-1, 1)
    W = torch.diag(pc_w)
    WA = torch.mm(W, A)
    ATWA = torch.mm(A.permute(1, 0), WA)

    Wb = torch.mm(W, b)
    ATWb = torch.mm(A.permute(1, 0), Wb)
    try:
        if torch.linalg.matrix_rank(ATWA) == 3:
            X = torch.linalg.solve(ATWA, ATWb)
        else:
            ATWA_1 = torch.pinverse(ATWA)
            X = torch.mm(ATWA_1, ATWb)
    except:
        # print('error when computing plane parameter due to ill-conditioned matrix')
        X = torch.ones((3, 1), device=A.device)
    return X
//  [python]
def test_inverse_covariance_matrices(self, dataset, Model):
        detector = self.train_detector(dataset, Model, MahalanobisDetector)
        assert detector.covariances.keys() == detector.inv_covariances.keys()
        for layer_name, cov in detector.covariances.items():
            inv_cov = detector.inv_covariances[layer_name]
            assert inv_cov.size() == cov.size()

            # Check that inverse is (pseudo) inverse
            rank = torch.matrix_rank(cov, rtol=self.rcond)
            assert torch.matrix_rank(inv_cov, rtol=self.rcond) == rank

            # TODO I'm uncertain which tolerances to use here, this is a
            # guesstimate based on some of the computations that are done and
            # test is still flaky, either computations are wrong or I skipped
            # too many steps
            assert torch.allclose(
                cov,
                cov @ inv_cov @ cov,
                rtol=(4 * cov.size(0) ** 2 * torch.finfo(cov.dtype).resolution),
                atol=(cov.size(0) * torch.finfo(cov.dtype).resolution),
            )
//  [python]
def test_inverse_covariance_matrices(self, dataset, Model):
        detector = self.train_detector(dataset, Model, MahalanobisDetector)
        assert detector.covariances.keys() == detector.inv_covariances.keys()
        for layer_name, cov in detector.covariances.items():
            inv_cov = detector.inv_covariances[layer_name]
            assert inv_cov.size() == cov.size()

            # Check that inverse is (pseudo) inverse
            rank = torch.linalg.matrix_rank(cov, rtol=self.rcond)
            assert torch.linalg.matrix_rank(inv_cov, rtol=self.rcond) == rank

            # TODO I'm uncertain which tolerances to use here, this is a
            # guesstimate based on some of the computations that are done and
            # test is still flaky, either computations are wrong or I skipped
            # too many steps
            assert torch.allclose(
                cov,
                cov @ inv_cov @ cov,
                rtol=(4 * cov.size(0) ** 2 * torch.finfo(cov.dtype).resolution),
                atol=(cov.size(0) * torch.finfo(cov.dtype).resolution),
            )
//  [python]
def get_plane(pc, pc_w):
    if torch.max(pc_w) < 1e-6:
        print('plane weight is too small!!')
        print(torch.max(pc_w))
        # return None, None, None

    pc_w = pc_w / torch.max(pc_w)
    # min least square
    n = pc.shape[0]
    A = torch.cat([pc[:, :2], torch.ones([n, 1], device=pc.device)], dim=-1)
    b = pc[:, 2].view(-1, 1)
    W = torch.diag(pc_w)
    WA = torch.mm(W, A)
    ATWA = torch.mm(A.permute(1, 0), WA)

    Wb = torch.mm(W, b)
    ATWb = torch.mm(A.permute(1, 0), Wb)
    try:
        if torch.linalg.matrix_rank(ATWA) == 3:
            X = torch.linalg.solve(ATWA, ATWb)
        else:
            ATWA_1 = torch.pinverse(ATWA)
            X = torch.mm(ATWA_1, ATWb)
    except:
        print('error when computing plane parameter due to ill-conditioned matrix')
        return None, None, None

    dn_up = torch.cat([X[0] * X[2], X[1] * X[2], -X[2]], dim=0),
    dn_norm = X[0] * X[0] + X[1] * X[1] + 1.0
    dn = dn_up[0] / dn_norm

    normal_n = dn / (torch.norm(dn) + 1e-10)
    for_p2plane = X[2] / torch.sqrt(dn_norm)
    return normal_n, dn, for_p2plane
//  [python]
def test_whitening_matrices(self, dataset, Model):
        detector = self.train_detector(dataset, Model, QuantumEntropyDetector)
        assert detector.covariances.keys() == detector.whitening_matrices.keys()
        for layer_name, cov in detector.covariances.items():
            W = detector.whitening_matrices[layer_name]
            assert W.size() == cov.size()

            # Check that Whitening matrix computes (pseudo) inverse
            rank = torch.matrix_rank(cov, rtol=self.rcond)
            assert torch.matrix_rank(W, rtol=self.rcond) == rank
            inv_cov = W @ W.mT
            assert torch.matrix_rank(inv_cov, rtol=self.rcond) == rank

            # TODO I'm uncertain which tolerances to use here, this is a
            # guesstimate based on some of the computations that are done and
            # test is still flaky, either computations are wrong or I skipped
            # too many steps
            assert torch.allclose(
                cov,
                cov @ inv_cov @ cov,
                rtol=(4 * cov.size(0) ** 2 * torch.finfo(cov.dtype).resolution),
                atol=(cov.size(0) * torch.finfo(cov.dtype).resolution),
            )
//  [python]
def test_whitening_matrices(self, dataset, Model):
        detector = self.train_detector(dataset, Model, QuantumEntropyDetector)
        assert detector.covariances.keys() == detector.whitening_matrices.keys()
        for layer_name, cov in detector.covariances.items():
            W = detector.whitening_matrices[layer_name]
            assert W.size() == cov.size()

            # Check that Whitening matrix computes (pseudo) inverse
            rank = torch.linalg.matrix_rank(cov, rtol=self.rcond)
            assert torch.linalg.matrix_rank(W, rtol=self.rcond) == rank
            inv_cov = W @ W.mT
            assert torch.linalg.matrix_rank(inv_cov, rtol=self.rcond) == rank

            # TODO I'm uncertain which tolerances to use here, this is a
            # guesstimate based on some of the computations that are done and
            # test is still flaky, either computations are wrong or I skipped
            # too many steps
            assert torch.allclose(
                cov,
                cov @ inv_cov @ cov,
                rtol=(4 * cov.size(0) ** 2 * torch.finfo(cov.dtype).resolution),
                atol=(cov.size(0) * torch.finfo(cov.dtype).resolution),
            )
//  [python]
def test_gevd(ch):
    stft = Stft(
        n_fft=8,
        win_length=None,
        hop_length=2,
        center=True,
        window="hann",
        normalized=False,
        onesided=True,
    )
    torch.random.manual_seed(0)
    x = random_speech[..., :ch]
    ilens = torch.LongTensor([16, 12])
    # (B, T, C, F) -> (B, F, C, T)
    X = torch.complex(*torch.unbind(stft(x, ilens)[0], dim=-1)).transpose(-1, -3)
    # (B, F, C, C)
    Phi_X = torch.einsum("...ct,...et->...ce", [X, X.conj()])

    is_singular = True
    while is_singular:
        N = torch.randn_like(X)
        Phi_N = torch.einsum("...ct,...et->...ce", [N, N.conj()])
        is_singular = not torch.linalg.matrix_rank(Phi_N).eq(ch).all()
    # Phi_N = torch.eye(ch, dtype=Phi_X.dtype).view(1, 1, ch, ch).expand_as(Phi_X)

    # e_val: (B, F, C)
    # e_vec: (B, F, C, C)
    e_val, e_vec = generalized_eigenvalue_decomposition(Phi_X, Phi_N)
    e_val = e_val.to(dtype=e_vec.dtype)
    assert torch.allclose(
        torch.matmul(Phi_X, e_vec),
        torch.matmul(torch.matmul(Phi_N, e_vec), e_val.diag_embed()),
    )
//  [python]
def test_rename_one():
    A = lo.IdentityOperator(3)

    assert (torch.matrix_rank(lo.operator_matrix_product(A, torch.eye(3)))) == 3

    cg = CG(A)
    b = torch.tensor([1., 4. , 6.])
    x = cg(b)

    assert torch.linalg.norm(b - A @ x) <= 1e-3
//  [python]
def test_rename_one():
    A = lo.IdentityOperator(3)

    assert (torch.linalg.matrix_rank(lo.operator_matrix_product(A, torch.eye(3)))) == 3

    cg = CG(A)
    b = torch.tensor([1., 4. , 6.])
    x = cg(b)

    assert torch.linalg.norm(b - A @ x) <= 1e-3
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self._device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.Reg_lambda * torch.eye(HTH.size(0))).to(self._device)
        A = HTH + I
        orig_rank = torch.matrix_rank(HTH).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            try:
                # W = torch.mm(torch.linalg.pinv(A.to(self.device)),
                #             HTY.to(self.device)).t().to(self._device)
                W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(self._device)
            except:
                W = torch.mm(torch.linalg.pinv(A.cpu()),
                            HTY.cpu()).t().to(self._device)
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self._device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.Reg_lambda * torch.eye(HTH.size(0))).to(self._device)
        A = HTH + I
        # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            try:
                # W = torch.mm(torch.linalg.pinv(A.to(self.device)),
                #             HTY.to(self.device)).t().to(self._device)
                W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(self._device)
            except:
                W = torch.mm(torch.linalg.pinv(A.cpu()),
                            HTY.cpu()).t().to(self._device)
        return W, tag
//  [python]
def laplacian_rank(H_list, device):
    # L = I - Dv^(-1/2) W De^(-1) H^(T) Dv^(-1/2)
    rank_list = []
    for tmpH in H_list:
        H = tmpH.clone()

        ## 删除空边
        n_edge = H.shape[1]
        tmp_sum = H.sum(dim=0)
        index = []
        for i in range(n_edge):
            if tmp_sum[i] != 0:
                index.append(i)

        H = H[:, index]
        ################
        n_node = H.shape[0]
        n_edge = H.shape[1]

        # the weight of the hyperedge
        # W = np.ones(n_edge)
        W = torch.ones(n_edge).to(device)

        # the degree of the node
        # DV = np.sum(H * W, axis=1)
        DV = torch.sum(H * W, axis=1)

        # the degree of the hyperedge
        # DE = np.sum(H, axis=0)
        DE = torch.sum(H, axis=0)

        # invDE = np.mat(np.diag(np.power(DE, -1)))
        invDE = torch.diag(torch.pow(DE,-1))

        # DV2 = np.mat(np.diag(np.power(DV, -0.5)))
        DV2 = torch.diag(torch.pow(DV, -0.5))

        # W = np.mat(np.diag(W))
        # H = np.mat(H)
        HT = H.T

        I = torch.eye(n_node, n_node).to(device)

        L = I - DV2 @ H @ W @ invDE @ HT @ DV2

        U, S, V = torch.linalg.svd(L)
        rank_L = torch.sum(S > 1e-5).item()

        rank_list.append(rank_L)

    print("===========================> Rank of L is: ", rank_list)
//  [python]
def laplacian_rank(H_list, device):
    # L = I - Dv^(-1/2) W De^(-1) H^(T) Dv^(-1/2)
    rank_list = []
    for tmpH in H_list:
        H = tmpH.clone()

        ## 删除空边
        n_edge = H.shape[1]
        tmp_sum = H.sum(dim=0)
        index = []
        for i in range(n_edge):
            if tmp_sum[i] != 0:
                index.append(i)

        H = H[:, index]
        ################
        n_node = H.shape[0]
        n_edge = H.shape[1]

        # the weight of the hyperedge
        # W = np.ones(n_edge)
        W = torch.ones(n_edge).to(device)

        # the degree of the node
        # DV = np.sum(H * W, axis=1)
        DV = torch.sum(H * W, axis=1)

        # the degree of the hyperedge
        # DE = np.sum(H, axis=0)
        DE = torch.sum(H, axis=0)

        # invDE = np.mat(np.diag(np.power(DE, -1)))
        invDE = torch.diag(torch.pow(DE,-1))

        # DV2 = np.mat(np.diag(np.power(DV, -0.5)))
        DV2 = torch.diag(torch.pow(DV, -0.5))

        # W = np.mat(np.diag(W))
        # H = np.mat(H)
        HT = H.T

        I = torch.eye(n_node, n_node).to(device)

        L = I - DV2 @ H @ W @ invDE @ HT @ DV2

        rank_L = torch.linalg.matrix_rank(L)

        rank_list.append(rank_L)

    print("===========================> Rank of L is: ", rank_list)
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self._device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.opts.lambda_reg * torch.eye(HTH.size(0))).to(self._device)
        A = HTH + I
        orig_rank = torch.matrix_rank(HTH).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self._device)
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self._device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (self.opts.lambda_reg * torch.eye(HTH.size(0))).to(self._device)
        A = HTH + I
        # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self._device)
        return W, tag
//  [python]
def diff_lstq(self, A, Y, lamb=0.010, log_file=None):
        """
        Differentiable least square
        param A: m x n
        param Y: n x 1
        """
        # Assuming A to be full column rank
        cols = A.shape[1]
        tmp_start = time.time()
        if cols == torch.linalg.matrix_rank(A):
            q, r = torch.linalg.qr(A)
            x = torch.inverse(r) @ q.T @ Y
        else:
            A_dash = A.permute(1, 0) @ A + lamb * torch.eye(cols)
            Y_dash = A.permute(1, 0) @ Y
            x = self.diff_lstq(A_dash, Y_dash)
        log_string(log_file, 'time of a single lstsq calc (seconds): {}'.format(time.time() - tmp_start))
        return x
//  [python]
def get_eff_dim(Z, t_uniformity=2):
    corr_coef = Z.T.corrcoef()
    # remove if all nan
    nan_cols = corr_coef.isnan().all(1)
    corr_coef = corr_coef[~nan_cols]
    corr_coef = corr_coef.T[~nan_cols].T
    rank= torch.linalg.matrix_rank(corr_coef, atol=1e-4, rtol=0.01, hermitian=True)
    rank_Z = torch.matrix_rank(Z, atol=1e-4, rtol=0.01)
    log_eigenv = torch.linalg.eigvalsh(corr_coef).abs().log().sort(descending=True)[0]

    # for memory reasons, uniformity has to be computes on batches (can't compute pairwise distance of 1M examples)
    uniformity = 0
    n_batches = 1000
    for batch in Z[torch.randperm(len(Z))].chunk(n_batches):
        # uniformity is computed on normalized features
        batch = torch.nn.functional.normalize(batch, dim=-1)
        uniformity += torch.pdist(batch, p=2).pow(2).mul(-t_uniformity).exp().mean()
    uniformity = (uniformity / n_batches).log()

    return rank_Z.cpu().numpy(), rank.cpu().numpy(), log_eigenv.cpu().numpy(), uniformity.cpu().numpy()
//  [python]
def get_eff_dim(Z, t_uniformity=2):
    corr_coef = Z.T.corrcoef()
    # remove if all nan
    nan_cols = corr_coef.isnan().all(1)
    corr_coef = corr_coef[~nan_cols]
    corr_coef = corr_coef.T[~nan_cols].T
    rank= torch.linalg.matrix_rank(corr_coef, atol=1e-4, rtol=0.01, hermitian=True)
    rank_Z = torch.linalg.matrix_rank(Z, atol=1e-4, rtol=0.01)
    log_eigenv = torch.linalg.eigvalsh(corr_coef).abs().log().sort(descending=True)[0]

    # for memory reasons, uniformity has to be computes on batches (can't compute pairwise distance of 1M examples)
    uniformity = 0
    n_batches = 1000
    for batch in Z[torch.randperm(len(Z))].chunk(n_batches):
        # uniformity is computed on normalized features
        batch = torch.nn.functional.normalize(batch, dim=-1)
        uniformity += torch.pdist(batch, p=2).pow(2).mul(-t_uniformity).exp().mean()
    uniformity = (uniformity / n_batches).log()

    return rank_Z.cpu().numpy(), rank.cpu().numpy(), log_eigenv.cpu().numpy(), uniformity.cpu().numpy()
//  [python]
def test_kernel_rank(kernel, data_x):
    for p in [5, 10, 20, 30]:
        # NOTE: single-precision runs into some numerical precision issues, but we can assert that the logic makes sense
        # using doubles.
        k_xx = kernel(data_x[:, :p].double())
        assert torch.matrix_rank(k_xx, hermitian=True) == min(size_m, kernel.test_rank(p))
//  [python]
def test_kernel_rank(kernel, data_x):
    for p in [5, 10, 20, 30]:
        # NOTE: single-precision runs into some numerical precision issues, but we can assert that the logic makes sense
        # using doubles.
        k_xx = kernel(data_x[:, :p].double())
        assert torch.linalg.matrix_rank(k_xx, hermitian=True) == min(size_m, kernel.test_rank(p))
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self._device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        # I = (self.opts.reg_lambda * torch.eye(HTH.size(0))).to(self._device)
        A = HTH
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            try:
                W = torch.linalg.lstsq(
                    A.cpu(), HTY.cpu(), driver='gelsd').solution.T.to(self._device)
            except:
                W = torch.mm(torch.linalg.pinv(A.cpu()),
                             HTY.cpu()).t().to(self._device)
        return W, tag
//  [python]
def test_make_random_matrix():
    A = make_random_matrix(50, 27)
    assert torch.matrix_rank(A) == 27
    B = make_random_matrix(50, 27, random_seed=200)
    with torch.no_grad():
        npt.assert_array_equal(B, make_random_matrix(50, 27, random_seed=200))
//  [python]
def test_make_random_matrix():
    A = make_random_matrix(50, 27)
    assert torch.linalg.matrix_rank(A) == 27
    B = make_random_matrix(50, 27, random_seed=200)
    with torch.no_grad():
        npt.assert_array_equal(B, make_random_matrix(50, 27, random_seed=200))
//  [python]
def nle_pca(img, patchsize=7, conf=(1-1e-6), itr=3):
    """ Translated from MATLAB code:
    Paper: Noise Level Estimation Using Weak Textured Patches Of
    A Single Noisy Image 
    http://www.ok.sc.e.titech.ac.jp/res/NLE/AWGNestimation.html
    """
    kh = torch.tensor([1/2,0,-1/2]).float().reshape(1,1,1,3)
    imgh = F.conv2d(img, kh)**2

    kv = kh.transpose(2,3)
    imgv = F.conv2d(img, kv)**2

    Dh = convmtx2(kh, patchsize, patchsize)
    Dv = convmtx2(kv, patchsize, patchsize)
    DD = Dh.T@Dh + Dv.T@Dv
    r = torch.linalg.matrix_rank(DD, hermitian=True)

    Dtr = torch.trace(DD)
    shape = r/2.0
    scale = 2.0 * Dtr / float(r)
    tau0 = gamma.ppf(conf, shape, scale=scale)

    nlevel = np.empty((img.shape[1],))
    th     = np.empty((img.shape[1],))
    num    = np.empty((img.shape[1],))

    for cha in range(1,img.shape[1]+1):
        X  = im2col(img[:,cha-1:cha], patchsize, patchsize)
        Xh = im2col(imgh[:,cha-1:cha], patchsize, patchsize-2)
        Xv = im2col(imgv[:,cha-1:cha], patchsize-2, patchsize)
        Xtr = Xh.sum(dim=0) + Xv.sum(dim=0)

        ###### nosie level estimation #####
        tau = np.inf
        if X.shape[1] < X.shape[0]:
            sig2 = 0
        else:
            cov = (X@(X.T))/(X.shape[1]-1)
            d = torch.linalg.eigvalsh(cov)
            sig2 = d[0]

        for i in range(2,itr):
        ##### weak texture selection #####
            tau = sig2 * tau0
            p = Xtr < tau
            Xtr = Xtr[p]
            X = X[:,p]

            ###### nosie level estimation #####
            if X.shape[1] < X.shape[0]:
                break
            cov = (X @ (X.T))/(X.shape[1]-1)
            d = torch.linalg.eigvalsh(cov)
            sig2 = d[0]

        nlevel[cha-1] = np.sqrt(sig2)
        th[cha-1]     = tau.item()
        num[cha-1]    = X.shape[1]
    if img.shape[1] == 1:
        return nlevel[0], th[0], num[0]
    return nlevel, th, num