jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def implicit_runge_kutta_stage_solution_trad(
        sol, deltat, time, rhs, butcher_tableau, stage_unknowns):
    a_coef, c_coef = butcher_tableau[0], butcher_tableau[2]
    nstages = a_coef.shape[0]
    ndof = stage_unknowns.shape[0]//nstages
    stage_rhs = []
    for ii in range(nstages):
        stage_time = time+c_coef[ii]*deltat
        stage_unknown = stage_unknowns[ii*ndof:(ii+1)*ndof]
        stage_sol = sol+stage_unknown
        srhs = rhs(stage_sol, stage_time)
        stage_rhs.append(srhs)
    stage_rhs = torch.cat(stage_rhs)
    tmp = torch.kron(a_coef, torch.eye(ndof))
    new_stage_unknowns = torch.matmul(tmp, deltat*stage_rhs)
    return new_stage_unknowns, stage_rhs
//  [python]
def implicit_runge_kutta_stage_solution_trad(
        sol, deltat, time, rhs, butcher_tableau, stage_unknowns):
    a_coef, c_coef = butcher_tableau[0], butcher_tableau[2]
    nstages = a_coef.shape[0]
    ndof = stage_unknowns.shape[0]//nstages
    stage_rhs = []
    for ii in range(nstages):
        stage_time = time+c_coef[ii]*deltat
        stage_unknown = stage_unknowns[ii*ndof:(ii+1)*ndof]
        stage_sol = sol+stage_unknown
        srhs = rhs(stage_sol, stage_time)
        stage_rhs.append(srhs)
    stage_rhs = torch.cat(stage_rhs)
    tmp = torch.kron(a_coef, torch.eye(ndof))
    new_stage_unknowns = torch.linalg.multi_dot((tmp, deltat*stage_rhs))
    return new_stage_unknowns, stage_rhs
//  [python]
def multidot(H,x):
    ret = torch.mm(torch.mm(torch.t(H), H), x)
    return ret
//  [python]
def multidot(H,x):
    ret = torch.linalg.multi_dot([torch.t(H), H, x]) 
    return ret
//  [python]
def test_shape_inference_unsupported_symbols_skip_fusion_group(self):
        """Test Glow shape inference unsupported symbols including skipping of
        symbols after a secondary fusion group."""

        def f(a, b):
            x1 = a * b
            x2 = x1 * b
            x3 = x2 * a
            x4 = x3 / b
            x5 = x4 / a
            x6 = x5 / b
            x7 = x6 * a
            x8 = x7 * b
            return x8 * torch.linalg.multi_dot([x8, x8])

        torch_glow.enableFusionPass_DO_NOT_USE_THIS()
        torch_glow.setFusionStartIndex(3)
        torch_glow.setFusionEndIndex(6)

        a = torch.randn(5, 5)
        b = torch.randn(5, 5)

        jit_f = torch.jit.trace(f, (a, b))

        jit_f_graph = jit_f.graph_for(a, b)

        torch_glow.clearFusionIndices()

        args = (a, b)

        # Don't skip nodes after the last fusion node.
        # in this case, one of the nodes (linalg.multi_dot) following the last fusion node
        # is not supported, and should be reported.
        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args, skip_last_fusion_node=False
        )
        expected = [
            "aten::linalg_multi_dot",
        ]
        self.assertEqual(set(expected), set(actual))

        # DO skip nodes after the last fusion node.
        # in this case, one of the nodes (linalg.multi_dot) following the last fusion node
        # is not supported, but is suppressed due to the skip_last_fusion_node flag.
        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args, skip_last_fusion_node=True
        )
        expected = []
        self.assertEqual(set(expected), set(actual))
//  [python]
def get_y_neg_log_likelihood(self, Y, X, N):
        """
        Compute latent negative log-likelihood Ly
        
        Parameters
        ----------

        Y : tensor(dtype)
            observation matrix

        X : tensor(dtype)
            latent state matrix

        N : int
            number of data points
        
        Return
        ------
        L_y = D/2*log(|K_y(X,X)|) + 1/2*trace(K_y^-1*Y*W_y^2*Y) - N*log(|W_y|)

        """
        K_y = self.get_y_kernel(X,X)
        U, info = torch.linalg.cholesky_ex(K_y, upper = True)
        U_inv = torch.inverse(U)
        Ky_inv = torch.matmul(U_inv,U_inv.t())
        log_det_K_y = 2 * torch.sum(torch.log(torch.diag(U)))

        W = torch.diag(torch.exp(self.y_log_lambdas))
        W2 = torch.diag(torch.exp(self.y_log_lambdas)**2)
        log_det_W = 2 * torch.sum(self.y_log_lambdas)

        return self.D / 2 * log_det_K_y + 1 / 2 * \
            torch.trace(torch.matmul(torch.matmul(torch.matmul(Ky_inv, Y), W2), Y.transpose(0,1))) \
            - N * log_det_W
//  [python]
def get_y_neg_log_likelihood(self, Y, X, N):
        """
        Compute latent negative log-likelihood Ly
        
        Parameters
        ----------

        Y : tensor(dtype)
            observation matrix

        X : tensor(dtype)
            latent state matrix

        N : int
            number of data points
        
        Return
        ------
        L_y = D/2*log(|K_y(X,X)|) + 1/2*trace(K_y^-1*Y*W_y^2*Y) - N*log(|W_y|)

        """
        K_y = self.get_y_kernel(X,X)
        U, info = torch.linalg.cholesky_ex(K_y, upper = True)
        U_inv = torch.inverse(U)
        Ky_inv = torch.matmul(U_inv,U_inv.t())
        log_det_K_y = 2 * torch.sum(torch.log(torch.diag(U)))

        W = torch.diag(torch.exp(self.y_log_lambdas))
        W2 = torch.diag(torch.exp(self.y_log_lambdas)**2)
        log_det_W = 2 * torch.sum(self.y_log_lambdas)

        return self.D / 2 * log_det_K_y + 1 / 2 * \
            torch.trace(torch.linalg.multi_dot([Ky_inv, Y, W2, Y.transpose(0,1)])) \
            - N * log_det_W
//  [python]
def _compute_grad_q(
    coords_x: torch.Tensor, grad_x: torch.Tensor, ic: tico.ic.IC
) -> torch.Tensor:
    """Project the gradient into internal coordinates."""
    b_matrix = ic.compute_b(coords_x)
    g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

    return torch.mm(torch.mm(g_matrix_inv, b_matrix), grad_x)
//  [python]
def _compute_grad_q(
    coords_x: torch.Tensor, grad_x: torch.Tensor, ic: tico.ic.IC
) -> torch.Tensor:
    """Project the gradient into internal coordinates."""
    b_matrix = ic.compute_b(coords_x)
    g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

    return torch.linalg.multi_dot([g_matrix_inv, b_matrix, grad_x])
//  [python]
def map_x_dynamics(self, Xstar, flg_noise = False):
        """
        Map Xstar to GP dynamics output

        Parameters
        ----------

        Xstar : tensor(dtype)
            input latent state matrix 

        flg_noise : boolean
            add noise to kernel matrix

        Return
        ------

        mean_Xout_pred : mean of Xout prediction

        diag_var_Xout_pred : variance of Xout prediction

        """
        Xin, Xout, _ = self.get_Xin_Xout_matrices()
        
        Kx_star = self.get_x_kernel(Xin, Xstar,False)
   
        mean_Xout_pred = torch.mm(torch.mm(Xout.t(), self.Kx_inv), Kx_star).t()
        diag_var_Xout_pred_common = self.get_x_diag_kernel(Xstar, flg_noise) - \
            torch.sum(torch.matmul(Kx_star.t(), self.Kx_inv) * Kx_star.t(), dim = 1)
        x_log_lambdas = torch.exp(self.x_log_lambdas)**-2
        diag_var_Xout_pred = diag_var_Xout_pred_common.unsqueeze(1) * x_log_lambdas.unsqueeze(0)

        return mean_Xout_pred, diag_var_Xout_pred
//  [python]
def map_x_dynamics(self, Xstar, flg_noise = False):
        """
        Map Xstar to GP dynamics output

        Parameters
        ----------

        Xstar : tensor(dtype)
            input latent state matrix 

        flg_noise : boolean
            add noise to kernel matrix

        Return
        ------

        mean_Xout_pred : mean of Xout prediction

        diag_var_Xout_pred : variance of Xout prediction

        """
        Xin, Xout, _ = self.get_Xin_Xout_matrices()
        
        Kx_star = self.get_x_kernel(Xin, Xstar,False)
   
        mean_Xout_pred = torch.linalg.multi_dot([Xout.t(), self.Kx_inv, Kx_star]).t()
        diag_var_Xout_pred_common = self.get_x_diag_kernel(Xstar, flg_noise) - \
            torch.sum(torch.matmul(Kx_star.t(), self.Kx_inv) * Kx_star.t(), dim = 1)
        x_log_lambdas = torch.exp(self.x_log_lambdas)**-2
        diag_var_Xout_pred = diag_var_Xout_pred_common.unsqueeze(1) * x_log_lambdas.unsqueeze(0)

        return mean_Xout_pred, diag_var_Xout_pred
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, N):
        """
        Compute dynamics map negative log-likelihood Lx
        
        Parameters
        ----------

        Xout : tensor(dtype)
            dynamics map output matrix

        Xin : tensor(dtype)
            dynamics map input matrix

        N : int
            number of data points
        
        Return
        ------
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - (N-dyn_back_step)*log(|W_x|)

        """
        
        K_x = self.get_x_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper = True)
        U_inv = torch.inverse(U)
        Kx_inv = torch.matmul(U_inv, U_inv.t())
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        W = torch.diag(torch.exp(self.x_log_lambdas))
        W2 = torch.diag(torch.exp(self.x_log_lambdas)**2)
        log_det_W = 2 * torch.sum(self.x_log_lambdas)

        return self.d / 2 * log_det_K_x + 1 / 2 * \
            torch.trace(torch.matmul(torch.matmul(torch.matmul(Kx_inv, Xout), W2), Xout.transpose(0,1))) \
            - Xin.shape[0] * log_det_W
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, N):
        """
        Compute dynamics map negative log-likelihood Lx
        
        Parameters
        ----------

        Xout : tensor(dtype)
            dynamics map output matrix

        Xin : tensor(dtype)
            dynamics map input matrix

        N : int
            number of data points
        
        Return
        ------
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - (N-dyn_back_step)*log(|W_x|)

        """
        
        K_x = self.get_x_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper = True)
        U_inv = torch.inverse(U)
        Kx_inv = torch.matmul(U_inv, U_inv.t())
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        W = torch.diag(torch.exp(self.x_log_lambdas))
        W2 = torch.diag(torch.exp(self.x_log_lambdas)**2)
        log_det_W = 2 * torch.sum(self.x_log_lambdas)

        return self.d / 2 * log_det_K_x + 1 / 2 * \
            torch.trace(torch.linalg.multi_dot([Kx_inv, Xout, W2, Xout.transpose(0,1)])) \
            - Xin.shape[0] * log_det_W
//  [python]
def map_x_to_y(self, Xstar, flg_noise = False):
        """
        Map Xstar to observation space: return mean and variance
        
        Parameters
        ----------

        Xstar : tensor(dtype)
            input latent state matrix 

        flg_noise : boolean
            add noise to prediction variance

        Return
        ------

        mean_Y_pred : mean of Y prediction

        diag_var_Y_pred : variance of Y prediction


        """

        Y_obs = self.get_Y()
        Y_obs = torch.tensor(Y_obs, dtype = self.dtype, device = self.device)

        Ky_star = self.get_y_kernel(self.X, Xstar,False)

        mean_Y_pred = torch.mm(torch.mm(Y_obs.t(), self.Ky_inv),Ky_star).t()
        diag_var_Y_pred_common = self.get_y_diag_kernel(Xstar, flg_noise) - \
            torch.sum(torch.matmul(Ky_star.t(), self.Ky_inv) * Ky_star.t(), dim = 1)
        y_log_lambdas = torch.exp(self.y_log_lambdas)**-2
        diag_var_Y_pred = diag_var_Y_pred_common.unsqueeze(1) * y_log_lambdas.unsqueeze(0)

        return mean_Y_pred + torch.tensor(self.meanY, dtype = self.dtype, device = self.device), diag_var_Y_pred
//  [python]
def map_x_to_y(self, Xstar, flg_noise = False):
        """
        Map Xstar to observation space: return mean and variance
        
        Parameters
        ----------

        Xstar : tensor(dtype)
            input latent state matrix 

        flg_noise : boolean
            add noise to prediction variance

        Return
        ------

        mean_Y_pred : mean of Y prediction

        diag_var_Y_pred : variance of Y prediction


        """

        Y_obs = self.get_Y()
        Y_obs = torch.tensor(Y_obs, dtype = self.dtype, device = self.device)

        Ky_star = self.get_y_kernel(self.X, Xstar,False)

        mean_Y_pred = torch.linalg.multi_dot([Y_obs.t(), self.Ky_inv,Ky_star]).t()
        diag_var_Y_pred_common = self.get_y_diag_kernel(Xstar, flg_noise) - \
            torch.sum(torch.matmul(Ky_star.t(), self.Ky_inv) * Ky_star.t(), dim = 1)
        y_log_lambdas = torch.exp(self.y_log_lambdas)**-2
        diag_var_Y_pred = diag_var_Y_pred_common.unsqueeze(1) * y_log_lambdas.unsqueeze(0)

        return mean_Y_pred + torch.tensor(self.meanY, dtype = self.dtype, device = self.device), diag_var_Y_pred
//  [python]
def _get_collocation_velocities(self, psols):
        subdomain_vels = []
        for jj, model in enumerate(
                self.pressure_solver._decomp._subdomain_models):
            mesh = model.physics.mesh
            idx = np.arange(mesh.mesh_pts.shape[1])
            flux_jac_vals = model.physics.flux_jac(idx)
            # only works because diffusivity field is
            # constant across each subdomain
            subdomain_vels.append(torch.hstack([
                -torch.matmul(flux_jac_vals[dd], psols[jj])[:, None]
                for dd in range(len(flux_jac_vals))]))
        return subdomain_vels
//  [python]
def _get_collocation_velocities(self, psols):
        subdomain_vels = []
        for jj, model in enumerate(
                self.pressure_solver._decomp._subdomain_models):
            mesh = model.physics.mesh
            idx = np.arange(mesh.mesh_pts.shape[1])
            flux_jac_vals = model.physics.flux_jac(idx)
            # only works because diffusivity field is
            # constant across each subdomain
            subdomain_vels.append(torch.hstack([
                -torch.linalg.multi_dot(
                    (flux_jac_vals[dd], psols[jj]))[:, None]
                for dd in range(len(flux_jac_vals))]))
        return subdomain_vels
//  [python]
def lstq(self, A, Y, lamb=0.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]
        if cols == torch.matrix_rank(A):
            q, r = torch.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.lstq(A_dash, Y_dash)
        return x
//  [python]
def lstq(self, A, Y, lamb=0.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]
        if cols == torch.linalg.matrix_rank(A):
            q, r = torch.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.lstq(A_dash, Y_dash)
        return x
//  [python]
def best_lambda(A):
    """
    Takes an under determined system and small lambda value,
    and comes up with lambda that makes the matrix A + lambda I
    invertible. Assuming A to be square matrix.
    """
    lamb = 1e-6
    cols = A.shape[0]

    for i in range(7):
        A_dash = A + lamb * torch.eye(cols, device=A.get_device())
        if cols == torch.matrix_rank(A_dash):
            # we achieved the required rank
            break
        else:
            # factor by which to increase the lambda. Choosing 10 for performance.
            lamb *= 10
    return lamb
//  [python]
def best_lambda(A):
    """
    Takes an under determined system and small lambda value,
    and comes up with lambda that makes the matrix A + lambda I
    invertible. Assuming A to be square matrix.
    """
    lamb = 1e-6
    cols = A.shape[0]

    for i in range(7):
        A_dash = A + lamb * torch.eye(cols, device=A.get_device())
        if cols == torch.linalg.matrix_rank(A_dash):
            # we achieved the required rank
            break
        else:
            # factor by which to increase the lambda. Choosing 10 for performance.
            lamb *= 10
    return lamb
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global entropy
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global entropy
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.linalg.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def get_feature_hook_densenet(self, input, output):
    global feature_result
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b-12,b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def get_feature_hook_densenet(self, input, output):
    global feature_result
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.linalg.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b-12,b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [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():
            import ipdb;
            ipdb.set_trace()

        # Assuming A to be full column rank
        if cols == torch.matrix_rank(A):
            # Full column rank
            q, r = torch.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 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():
            import ipdb;
            ipdb.set_trace()

        # Assuming A to be full column rank
        if cols == torch.linalg.matrix_rank(A):
            # Full column rank
            q, r = torch.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 get_feature_hook_googlenet(self, input, output):
    global feature_result
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b-12,b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def get_feature_hook_googlenet(self, input, output):
    global feature_result
    global total
    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.linalg.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b-12,b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def check_full_rank(input):
    """check input matrix is full rank. if not, print rank."""
    if input.dim() < 3:
        input = input.unsqueeze(0)
    C = input.size(1)
    for mat in input:
        rank = torch.matrix_rank(mat)
        if rank < C:
            print(f'Rank: {rank}/{C}')
            return
//  [python]
def check_full_rank(input):
    """check input matrix is full rank. if not, print rank."""
    if input.dim() < 3:
        input = input.unsqueeze(0)
    C = input.size(1)
    for mat in input:
        rank = torch.linalg.matrix_rank(mat)
        if rank < C:
            print(f'Rank: {rank}/{C}')
            return
//  [python]
def compute_rank(tensor):
    # Compute the matrix rank of a tensor
    rank = torch.matrix_rank(tensor)
    return rank
//  [python]
def compute_rank(tensor):
    # Compute the matrix rank of a tensor
    rank = torch.linalg.matrix_rank(tensor)
    return rank
//  [python]
def lstq(Y, A, lamb=0.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]
    if cols == torch.matrix_rank(A):
        q, r = torch.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
        #if Y_dash.dim() == 1:
        #  Y_dash = Y_dash.view(-1, 1)
        x = lstq(Y_dash, A_dash)
    return x
//  [python]
def lstq(Y, A, lamb=0.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]
    if cols == torch.linalg.matrix_rank(A):
        q, r = torch.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
        #if Y_dash.dim() == 1:
        #  Y_dash = Y_dash.view(-1, 1)
        x = lstq(Y_dash, A_dash)
    return x
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global entropy
    global total #
    a = output.shape[0] # B 128
    b = output.shape[1] # C 16
    # 输出feature的秩（输入图片计算每一层的 平均秩的信息）>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    c = torch.tensor([torch.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b)]) #(2408, )

    c = c.view(a, -1).float() # (B, C)(128, 16)
    c = c.sum(0) # (C,)
    feature_result = feature_result * total + c
    total = total + a # B add
    feature_result = feature_result / total
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global entropy
    global total #
    a = output.shape[0] # B 128
    b = output.shape[1] # C 16
    # 输出feature的秩（输入图片计算每一层的 平均秩的信息）>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    c = torch.tensor([torch.linalg.matrix_rank(output[i,j,:,:]).item() for i in range(a) for j in range(b)]) #(2408, )

    c = c.view(a, -1).float() # (B, C)(128, 16)
    c = c.sum(0) # (C,)
    feature_result = feature_result * total + c
    total = total + a # B add
    feature_result = feature_result / total
//  [python]
def forward(self, x1, x2):

        z1 = self.encoder(x1)
        z2 = self.encoder(x2)
        corr_in = corrcoef(z1).abs().fill_diagonal_(0).mean()
        # rank_c = torch.matrix_rank(cov.cpu().detach())
        # rank_z_in = torch.matrix_rank(z1.detach().t().cpu())
        rank_c_in = torch.matrix_rank(covariance(z1.detach()).cpu().detach())
        z1 = self.wn(z1)
        z2 = self.wn(z2)
        corr_out = corrcoef(z1).abs().fill_diagonal_(0).mean()
        rank_c_out = torch.matrix_rank(covariance(z1.detach()).cpu().detach())
        # L = - F.cosine_similarity(z1, z2, dim=-1).mean()
        L = self.loss_fn(z1, z2)

        return {'loss': L,
                'corr_in': corr_in,
                'corr_out': corr_out, 'rank_in':rank_c_in, 'rank_out':rank_c_out}
//  [python]
def forward(self, x1, x2):

        z1 = self.encoder(x1)
        z2 = self.encoder(x2)
        corr_in = corrcoef(z1).abs().fill_diagonal_(0).mean()
        # rank_c = torch.linalg.matrix_rank(cov.cpu().detach())
        # rank_z_in = torch.linalg.matrix_rank(z1.detach().t().cpu())
        rank_c_in = torch.linalg.matrix_rank(covariance(z1.detach()).cpu().detach())
        z1 = self.wn(z1)
        z2 = self.wn(z2)
        corr_out = corrcoef(z1).abs().fill_diagonal_(0).mean()
        rank_c_out = torch.linalg.matrix_rank(covariance(z1.detach()).cpu().detach())
        # L = - F.cosine_similarity(z1, z2, dim=-1).mean()
        L = self.loss_fn(z1, z2)

        return {'loss': L,
                'corr_in': corr_in,
                'corr_out': corr_out, 'rank_in':rank_c_in, 'rank_out':rank_c_out}
//  [python]
def best_lambda(A):
    """
    Takes an under determined system and small lambda value,
    and comes up with lambda that makes the matrix A + lambda I
    invertible. Assuming A to be square matrix.
    """
    lamb = 1e-6
    cols = A.shape[0]

    for i in range(7):
        A_dash = A + lamb * torch.eye(cols)
        if cols == torch.matrix_rank(A_dash):
            # we achieved the required rank
            break
        else:
            # factor by which to increase the lambda. Choosing 10 for performance.
            lamb *= 10
    return lamb
//  [python]
def best_lambda(A):
    """
    Takes an under determined system and small lambda value,
    and comes up with lambda that makes the matrix A + lambda I
    invertible. Assuming A to be square matrix.
    """
    lamb = 1e-6
    cols = A.shape[0]

    for i in range(7):
        A_dash = A + lamb * torch.eye(cols)
        if cols == torch.linalg.matrix_rank(A_dash):
            # we achieved the required rank
            break
        else:
            # factor by which to increase the lambda. Choosing 10 for performance.
            lamb *= 10
    return lamb
//  [python]
def matrix_rank(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        result = torch.matrix_rank(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        ranks = []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            r = torch.matrix_rank(slices[i])
            # interesting,
            # ranks.append(r)
            # does not work on pytorch 1.0.0
            # but the below code does
            ranks += [r]
        result = torch.stack(ranks).view(batches)
    return result
//  [python]
def matrix_rank(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        result = torch.linalg.matrix_rank(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        ranks = []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            r = torch.linalg.matrix_rank(slices[i])
            # interesting,
            # ranks.append(r)
            # does not work on pytorch 1.0.0
            # but the below code does
            ranks += [r]
        result = torch.stack(ranks).view(batches)
    return result
//  [python]
def test_old_matrix_rank(self, device, dtype):
        a = torch.eye(10, dtype=dtype, device=device)
        self.assertEqual(torch.matrix_rank(a).item(), 10)
        self.assertEqual(torch.matrix_rank(a, True).item(), 10)

        a[5, 5] = 0
        self.assertEqual(torch.matrix_rank(a).item(), 9)
        self.assertEqual(torch.matrix_rank(a, True).item(), 9)

        a = torch.randn(24, 42, dtype=dtype, device=device)
        self.assertEqual(torch.matrix_rank(a), torch.matrix_rank(a.t()))
        aaT = torch.mm(a, a.conj().t())
        self.assertEqual(torch.matrix_rank(aaT), torch.matrix_rank(aaT, True))
        aTa = torch.mm(a.conj().t(), a)
        self.assertEqual(torch.matrix_rank(aTa), torch.matrix_rank(aTa, True))

        a = torch.randn(35, 75, dtype=dtype, device=device)
        self.assertEqual(torch.matrix_rank(a), np.linalg.matrix_rank(a.cpu().numpy()))
        self.assertEqual(torch.matrix_rank(a, 0.01), np.linalg.matrix_rank(a.cpu().numpy(), 0.01))

        aaT = torch.mm(a, a.conj().t())
        self.assertEqual(torch.matrix_rank(aaT), np.linalg.matrix_rank(aaT.cpu().numpy()))
        self.assertEqual(torch.matrix_rank(aaT, 0.01), np.linalg.matrix_rank(aaT.cpu().numpy(), 0.01))

        if np.lib.NumpyVersion(np.__version__) >= '1.14.0':
            self.assertEqual(torch.matrix_rank(aaT, True), np.linalg.matrix_rank(aaT.cpu().numpy(), True))
            self.assertEqual(torch.matrix_rank(aaT, 0.01, True), np.linalg.matrix_rank(aaT.cpu().numpy(), 0.01, True))
//  [python]
def test_old_matrix_rank(self, device, dtype):
        a = torch.eye(10, dtype=dtype, device=device)
        self.assertEqual(torch.linalg.matrix_rank(a).item(), 10)
        self.assertEqual(torch.linalg.matrix_rank(a, True).item(), 10)

        a[5, 5] = 0
        self.assertEqual(torch.linalg.matrix_rank(a).item(), 9)
        self.assertEqual(torch.linalg.matrix_rank(a, True).item(), 9)

        a = torch.randn(24, 42, dtype=dtype, device=device)
        self.assertEqual(torch.linalg.matrix_rank(a), torch.linalg.matrix_rank(a.t()))
        aaT = torch.mm(a, a.conj().t())
        self.assertEqual(torch.linalg.matrix_rank(aaT), torch.linalg.matrix_rank(aaT, True))
        aTa = torch.mm(a.conj().t(), a)
        self.assertEqual(torch.linalg.matrix_rank(aTa), torch.linalg.matrix_rank(aTa, True))

        a = torch.randn(35, 75, dtype=dtype, device=device)
        self.assertEqual(torch.linalg.matrix_rank(a), np.linalg.matrix_rank(a.cpu().numpy()))
        self.assertEqual(torch.linalg.matrix_rank(a, 0.01), np.linalg.matrix_rank(a.cpu().numpy(), 0.01))

        aaT = torch.mm(a, a.conj().t())
        self.assertEqual(torch.linalg.matrix_rank(aaT), np.linalg.matrix_rank(aaT.cpu().numpy()))
        self.assertEqual(torch.linalg.matrix_rank(aaT, 0.01), np.linalg.matrix_rank(aaT.cpu().numpy(), 0.01))

        if np.lib.NumpyVersion(np.__version__) >= '1.14.0':
            self.assertEqual(torch.linalg.matrix_rank(aaT, True), np.linalg.matrix_rank(aaT.cpu().numpy(), True))
            self.assertEqual(torch.linalg.matrix_rank(aaT, 0.01, True), np.linalg.matrix_rank(aaT.cpu().numpy(), 0.01, True))
//  [python]
def test_matrix_rank_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.matrix_rank(a)
//  [python]
def test_matrix_rank_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.matrix_rank(a)
//  [python]
def lstq(A, Y, lamb=0.01):
        """
        Differentiable least square
        :param A: m x n
        :param Y: n x 1
        """
        # Assuming A to be full column rank
        cols = A.shape[1]
        print (torch.matrix_rank(A))
        if cols == torch.matrix_rank(A):
            q, r = torch.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 = lstq(A_dash, Y_dash)
        return x
//  [python]
def lstq(A, Y, lamb=0.01):
        """
        Differentiable least square
        :param A: m x n
        :param Y: n x 1
        """
        # Assuming A to be full column rank
        cols = A.shape[1]
        print (torch.linalg.matrix_rank(A))
        if cols == torch.linalg.matrix_rank(A):
            q, r = torch.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 = lstq(A_dash, Y_dash)
        return x
//  [python]
def rank_torch(A, args=None):
    A_tensor = torch.from_numpy(A).cuda()    
    return torch.matrix_rank(A_tensor).cpu().item()
//  [python]
def rank_torch(A, args=None):
    A_tensor = torch.from_numpy(A).cuda()    
    return torch.linalg.matrix_rank(A_tensor).cpu().item()
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global total

    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.matrix_rank(output[i, j, :, :]).item() for i in range(a) for j in range(b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [python]
def get_feature_hook(self, input, output):
    global feature_result
    global total

    a = output.shape[0]
    b = output.shape[1]
    c = torch.tensor([torch.linalg.matrix_rank(output[i, j, :, :]).item() for i in range(a) for j in range(b)])

    c = c.view(a, -1).float()
    c = c.sum(0)
    feature_result = feature_result * total + c
    total = total + a
    feature_result = feature_result / total
//  [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():
            import ipdb; ipdb.set_trace()

        # Assuming A to be full column rank
        if cols == torch.matrix_rank(A):
            # Full column rank
            q, r = torch.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.
            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 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():
            import ipdb; ipdb.set_trace()

        # Assuming A to be full column rank
        if cols == torch.linalg.matrix_rank(A):
            # Full column rank
            q, r = torch.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.
            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 test_matrix_rank(self, device):
        a = torch.eye(10, device=device)
        self.assertEqual(torch.matrix_rank(a).item(), 10)
        self.assertEqual(torch.matrix_rank(a, True).item(), 10)

        a[5, 5] = 0
        self.assertEqual(torch.matrix_rank(a).item(), 9)
        self.assertEqual(torch.matrix_rank(a, True).item(), 9)

        a = torch.randn(24, 42, device=device)
        self.assertEqual(torch.matrix_rank(a), torch.matrix_rank(a.t()))
        aaT = torch.mm(a, a.t())
        self.assertEqual(torch.matrix_rank(aaT), torch.matrix_rank(aaT, True))
        aTa = torch.mm(a.t(), a)
        self.assertEqual(torch.matrix_rank(aTa), torch.matrix_rank(aTa, True))

        if TEST_NUMPY:
            from numpy.linalg import matrix_rank
            a = torch.randn(35, 75, device=device)
            self.assertEqual(torch.matrix_rank(a).item(), matrix_rank(a.cpu().numpy()))
            self.assertEqual(torch.matrix_rank(a, 0.01).item(), matrix_rank(a.cpu().numpy(), 0.01))

            aaT = torch.mm(a, a.t())
            self.assertEqual(torch.matrix_rank(aaT).item(), matrix_rank(aaT.cpu().numpy()))
            self.assertEqual(torch.matrix_rank(aaT, 0.01).item(), matrix_rank(aaT.cpu().numpy(), 0.01))

            if np.lib.NumpyVersion(np.__version__) >= '1.14.0':
                self.assertEqual(torch.matrix_rank(aaT, True).item(), matrix_rank(aaT.cpu().numpy(), True))
                self.assertEqual(torch.matrix_rank(aaT, 0.01, True).item(),
                                 matrix_rank(aaT.cpu().numpy(), 0.01, True))
//  [python]
def test_matrix_rank(self, device):
        a = torch.eye(10, device=device)
        self.assertEqual(torch.linalg.matrix_rank(a).item(), 10)
        self.assertEqual(torch.linalg.matrix_rank(a, True).item(), 10)

        a[5, 5] = 0
        self.assertEqual(torch.linalg.matrix_rank(a).item(), 9)
        self.assertEqual(torch.linalg.matrix_rank(a, True).item(), 9)

        a = torch.randn(24, 42, device=device)
        self.assertEqual(torch.linalg.matrix_rank(a), torch.linalg.matrix_rank(a.t()))
        aaT = torch.mm(a, a.t())
        self.assertEqual(torch.linalg.matrix_rank(aaT), torch.linalg.matrix_rank(aaT, True))
        aTa = torch.mm(a.t(), a)
        self.assertEqual(torch.linalg.matrix_rank(aTa), torch.linalg.matrix_rank(aTa, True))

        if TEST_NUMPY:
            from numpy.linalg import matrix_rank
            a = torch.randn(35, 75, device=device)
            self.assertEqual(torch.linalg.matrix_rank(a).item(), matrix_rank(a.cpu().numpy()))
            self.assertEqual(torch.linalg.matrix_rank(a, 0.01).item(), matrix_rank(a.cpu().numpy(), 0.01))

            aaT = torch.mm(a, a.t())
            self.assertEqual(torch.linalg.matrix_rank(aaT).item(), matrix_rank(aaT.cpu().numpy()))
            self.assertEqual(torch.linalg.matrix_rank(aaT, 0.01).item(), matrix_rank(aaT.cpu().numpy(), 0.01))

            if np.lib.NumpyVersion(np.__version__) >= '1.14.0':
                self.assertEqual(torch.linalg.matrix_rank(aaT, True).item(), matrix_rank(aaT.cpu().numpy(), True))
                self.assertEqual(torch.linalg.matrix_rank(aaT, 0.01, True).item(),
                                 matrix_rank(aaT.cpu().numpy(), 0.01, True))
//  [python]
def command(self, state, choose_best=False):
        if not torch.is_tensor(state):
            state = torch.tensor(state)
        state = state.to(dtype=self.dtype, device=self.d)

        self.reset()

        for m in range(self.M):
            top_samples = self._sample_top_trajectories(state, self.num_elite)
            # fit the gaussian to those samples
            self.mean = torch.mean(top_samples, dim=0)
            self.cov = pytorch_cov(top_samples, rowvar=False)
            if torch.matrix_rank(self.cov) < self.cov.shape[0]:
                self.cov += self.cov_reg

        if choose_best and self.choose_best:
            top_sample = self._sample_top_trajectories(state, 1)
        else:
            top_sample = self.action_distribution.sample((1,))

        # only apply the first action from this trajectory
        u = top_sample[0, self._slice_control(0)]

        return u
//  [python]
def command(self, state, choose_best=False):
        if not torch.is_tensor(state):
            state = torch.tensor(state)
        state = state.to(dtype=self.dtype, device=self.d)

        self.reset()

        for m in range(self.M):
            top_samples = self._sample_top_trajectories(state, self.num_elite)
            # fit the gaussian to those samples
            self.mean = torch.mean(top_samples, dim=0)
            self.cov = pytorch_cov(top_samples, rowvar=False)
            if torch.linalg.matrix_rank(self.cov) < self.cov.shape[0]:
                self.cov += self.cov_reg

        if choose_best and self.choose_best:
            top_sample = self._sample_top_trajectories(state, 1)
        else:
            top_sample = self.action_distribution.sample((1,))

        # only apply the first action from this trajectory
        u = top_sample[0, self._slice_control(0)]

        return u
//  [python]
def forward(self, x1, x2, labels=None):
        if self.get_feature:

            f1 = self.backbone(x1)
            f2 = self.backbone(x2)
            z1 = self.projector(f1)
            z2 = self.projector(f2)

        else:
            z1 = self.encoder(x1)
            z2 = self.encoder(x2)
        N, D = z1.shape

        z1_norm = (z1 - z1.mean(dim=0, keepdim=True)) / z1.std(dim=0, keepdim=True)
        z2_norm = (z2 - z2.mean(dim=0, keepdim=True)) / z2.std(dim=0, keepdim=True)

        c = torch.mm(z1_norm.t(), z2_norm) / N
        # breakpoint()
        c_diff = (c - torch.eye(D).to(c)).pow(2)
        # print(c_diff[:5,:5])
        c_diff.flatten()[1:].view(D-1, D+1)[:,:-1].mul_(self.lamb)
        # off_diagonal(c_diff)
        # print(c_diff[:5,:5])
        # breakpoint()
        loss = c_diff.sum()
        if self.get_feature:
            return {'loss':loss, 'rank':torch.tensor(0), 'corr':torch.tensor(0),  'feature1':[f1.detach(), f2.detach()], 'feature2':[z1.detach(), z2.detach()]}
        else:
            return {'loss':loss, 'rank':torch.matrix_rank(z1.detach().cpu()), 'corr': corrcoef(z1.detach()).abs().fill_diagonal_(0).sum() / (D*(D-1))}
//  [python]
def forward(self, x1, x2, labels=None):
        if self.get_feature:

            f1 = self.backbone(x1)
            f2 = self.backbone(x2)
            z1 = self.projector(f1)
            z2 = self.projector(f2)

        else:
            z1 = self.encoder(x1)
            z2 = self.encoder(x2)
        N, D = z1.shape

        z1_norm = (z1 - z1.mean(dim=0, keepdim=True)) / z1.std(dim=0, keepdim=True)
        z2_norm = (z2 - z2.mean(dim=0, keepdim=True)) / z2.std(dim=0, keepdim=True)

        c = torch.mm(z1_norm.t(), z2_norm) / N
        # breakpoint()
        c_diff = (c - torch.eye(D).to(c)).pow(2)
        # print(c_diff[:5,:5])
        c_diff.flatten()[1:].view(D-1, D+1)[:,:-1].mul_(self.lamb)
        # off_diagonal(c_diff)
        # print(c_diff[:5,:5])
        # breakpoint()
        loss = c_diff.sum()
        if self.get_feature:
            return {'loss':loss, 'rank':torch.tensor(0), 'corr':torch.tensor(0),  'feature1':[f1.detach(), f2.detach()], 'feature2':[z1.detach(), z2.detach()]}
        else:
            return {'loss':loss, 'rank':torch.linalg.matrix_rank(z1.detach().cpu()), 'corr': corrcoef(z1.detach()).abs().fill_diagonal_(0).sum() / (D*(D-1))}
//  [python]
def register(self, model):
        def hook(module, input, output):
            side = self.side
            x = input[0] if side == 'input' else output

            path = self.get_rank_path(model, module)
            if not path.exists():
                # compute and store ranks
                start = time.time()
                mean = torch.matrix_rank(x).float().mean(dim=0)
                end = time.time()
                setattr(module, self.name, mean)
                torch.save(mean, path)

                logging.info(f"Computed rank for {module._name} ({x.shape}) in {round(end - start, 2)}s")
            else:
                logging.debug(f"Loaded precomputed rank for {module._name})")

        model._handles = []
        model._name = ''
        is_hook_registered = False
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                module._name = name
                path = self.get_rank_path(model, module)
                if not path.exists():
                    model._handles.append(module.register_forward_hook(hook))
                    is_hook_registered = True
        return is_hook_registered
//  [python]
def register(self, model):
        def hook(module, input, output):
            side = self.side
            x = input[0] if side == 'input' else output

            path = self.get_rank_path(model, module)
            if not path.exists():
                # compute and store ranks
                start = time.time()
                mean = torch.linalg.matrix_rank(x).float().mean(dim=0)
                end = time.time()
                setattr(module, self.name, mean)
                torch.save(mean, path)

                logging.info(f"Computed rank for {module._name} ({x.shape}) in {round(end - start, 2)}s")
            else:
                logging.debug(f"Loaded precomputed rank for {module._name})")

        model._handles = []
        model._name = ''
        is_hook_registered = False
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                module._name = name
                path = self.get_rank_path(model, module)
                if not path.exists():
                    model._handles.append(module.register_forward_hook(hook))
                    is_hook_registered = True
        return is_hook_registered
//  [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 + 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.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:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self.device)
        return W, tag
//  [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.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