jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def calculate_composition_weights(
    datasets: Union[Dataset, List[Dataset]], property: str
) -> Tuple[torch.Tensor, List[int]]:
    """Calculate the composition weights for a dataset.

    For now, it assumes per-system properties.

    :param dataset: Dataset to calculate the composition weights for.
    :returns: Composition weights for the dataset, as well as the
        list of species that the weights correspond to.
    """
    if not isinstance(datasets, list):
        datasets = [datasets]

    species = get_all_species(datasets)
    # note that this is sorted, and the composition weights are sorted
    # as well, because the species are sorted in the composition features

    targets = torch.stack(
        [
            sample._asdict()[property].block().values
            for dataset in datasets
            for sample in dataset
        ]
    )
    targets = targets.squeeze(dim=(1, 2))  # remove component and property dimensions

    structure_list = [
        sample._asdict()["system"] for dataset in datasets for sample in dataset
    ]

    dtype = structure_list[0].positions.dtype
    composition_features = torch.empty((len(structure_list), len(species)), dtype=dtype)
    for i, structure in enumerate(structure_list):
        for j, s in enumerate(species):
            composition_features[i, j] = torch.sum(structure.types == s)

    regularizer = 1e-20
    while regularizer:
        if regularizer > 1e5:
            raise RuntimeError(
                "Failed to solve the linear system to calculate the "
                "composition weights. The dataset is probably too small "
                "or ill-conditioned."
            )
        try:
            solution = torch.linalg.solve(
                composition_features.T @ composition_features
                + regularizer
                * torch.eye(
                    composition_features.shape[1],
                    dtype=composition_features.dtype,
                    device=composition_features.device,
                ),
                composition_features.T @ targets,
            )
            break
        except torch._C._LinAlgError:
            regularizer *= 10.0

    return solution, species
//  [python]
def differentiate_dloss_db(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.solve(b.unsqueeze(1), K)[0].squeeze()
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()

    db_matrix = torch.inverse(K.T) @ xhat.grad
    db_matrix_free = torch.solve(xhat.grad.unsqueeze(1), K.T)[0].squeeze()
    
    return db_matrix, db_matrix_free
//  [python]
def differentiate_dloss_db(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.linalg.solve(K, b)
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()

    db_matrix = torch.inverse(K.T) @ xhat.grad
    db_matrix_free = torch.linalg.solve(K.T, xhat.grad)
    
    return db_matrix, db_matrix_free
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = self.B  # (H N)
        P = self.P
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.copies)
        P = repeat(P, "r t n -> r (v t) n", v=self.copies)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.copies)
        w = repeat(w, "t n -> (v t) n", v=self.copies)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H r r)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.solve(Q_D, R.to(Q_D))[0]  # (H r N)
        except RuntimeError:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).cpu(), Q_D.cpu())).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (r H N)
            "P": P,  # (r H N)
            "Q": Q,  # (r H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = self.B  # (H N)
        P = self.P
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.copies)
        P = repeat(P, "r t n -> r (v t) n", v=self.copies)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.copies)
        w = repeat(w, "t n -> (v t) n", v=self.copies)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H r r)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.linalg.solve(R.to(Q_D), Q_D)  # (H r N)
        except torch._C._LinAlgError:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).cpu(), Q_D.cpu())).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (r H N)
            "P": P,  # (r H N)
            "Q": Q,  # (r H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping"""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C))  # (H C N)

        dC = torch.solve(
            C.unsqueeze(-1),
            I - dA_L.transpose(-1, -2),
        ).solution.squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == "linear":
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2 * self.dC[:, :, : self.N]
        elif mode == "diagonal":
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print(
                    "Diagonalization error:",
                    torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA),
                )

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract("h n m, h m -> h n", V_inv, self.dB)
            self.dC = contract("h n m, c h n -> c h m", V, self.dC)

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def setup_step(self, mode="dense"):
        """Set up dA, dB, dC discretized parameters for stepping"""
        self.dA, self.dB = self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C))  # (H C N)

        dC = torch.linalg.solve(
            I - dA_L.transpose(-1, -2),
            C.unsqueeze(-1),
        ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == "linear":
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2 * self.dC[:, :, : self.N]
        elif mode == "diagonal":
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print(
                    "Diagonalization error:",
                    torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA),
                )

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract("h n m, h m -> h n", V_inv, self.dB)
            self.dC = contract("h n m, c h n -> c h m", V, self.dC)

        elif mode == "dense":
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
            )
//  [python]
def matrix_differentiate_dloss_dtheta(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.solve(b.unsqueeze(1), K)[0].squeeze(1)
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()
    
    
    def Kmat(theta):
        return theta * 2

    dK_dtheta = torch.autograd.functional.jacobian(Kmat, theta)
    dxhat_dtheta = -torch.inverse(K) @ dK_dtheta @ xhat

    dloss_dtheta = dxhat_dtheta @ xhat.grad

    return dloss_dtheta
//  [python]
def matrix_differentiate_dloss_dtheta(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.linalg.solve(K, b)
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()
    
    
    def Kmat(theta):
        return theta * 2

    dK_dtheta = torch.autograd.functional.jacobian(Kmat, theta)
    dxhat_dtheta = -torch.inverse(K) @ dK_dtheta @ xhat

    dloss_dtheta = dxhat_dtheta @ xhat.grad

    return dloss_dtheta
//  [python]
def matrix_free_differentiate_dloss_dtheta(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.solve(b.unsqueeze(1), K)[0].squeeze(1)
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()
    
    def linop(theta):
        return theta*2 @ xhat

    db_dtheta = -torch.autograd.functional.jacobian(linop, theta).permute(1,2,0).unsqueeze(-1)
    dxhat_dtheta = torch.solve(db_dtheta, K)[0].squeeze(-1)

    dloss_dtheta = dxhat_dtheta @ xhat.grad

    return dloss_dtheta
//  [python]
def matrix_free_differentiate_dloss_dtheta(seed):
    torch.manual_seed(seed)
    theta = torch.randn((32,32), requires_grad=True)
    K = theta * 2
    x = torch.randn((32))
    b = K @ x
    b = b.clone().detach().requires_grad_(True)

    xhat = torch.linalg.solve(K, b)
    xhat.retain_grad()

    loss = xhat.mean()
    loss.backward()
    
    def linop(theta):
        return theta*2 @ xhat

    db_dtheta = -torch.autograd.functional.jacobian(linop, theta).permute(1,2,0).unsqueeze(-1)
    dxhat_dtheta = torch.linalg.solve(K, db_dtheta).squeeze(-1)

    dloss_dtheta = dxhat_dtheta @ xhat.grad

    return dloss_dtheta
//  [python]
def test_solve_lyapunov_diag():
    d = 50
    # setting up some matrices
    Q = torch.randn(d,).diag()
    K = torch.randn(d,).diag().pow(2) + 1e-3 * torch.eye(d)
    # solving the lyapunov equation using diag
    st_time = timeit.default_timer()
    X = torch.diag(solve_lyapunov_diag(K.diag(), Q.diag()))
    fast_time = timeit.default_timer() - st_time
    Q_solve = K @ X.t() + X @ K.t()
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q, decimal=4)
    # check it's actually fast
    st_time = timeit.default_timer()
    # naive solve
    I = torch.eye(d)
    LHS = torch.kron(I, K) + torch.kron(K, I)
    vecX = torch.solve(Q.t().ravel().unsqueeze(1), LHS)[0].ravel()
    slow_time = timeit.default_timer() - st_time
    assert slow_time > fast_time
//  [python]
def test_solve_lyapunov_diag():
    d = 50
    # setting up some matrices
    Q = torch.randn(d,).diag()
    K = torch.randn(d,).diag().pow(2) + 1e-3 * torch.eye(d)
    # solving the lyapunov equation using diag
    st_time = timeit.default_timer()
    X = torch.diag(solve_lyapunov_diag(K.diag(), Q.diag()))
    fast_time = timeit.default_timer() - st_time
    Q_solve = K @ X.t() + X @ K.t()
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q, decimal=4)
    # check it's actually fast
    st_time = timeit.default_timer()
    # naive solve
    I = torch.eye(d)
    LHS = torch.kron(I, K) + torch.kron(K, I)
    vecX = torch.linalg.solve(LHS, Q.t().ravel())
    slow_time = timeit.default_timer() - st_time
    assert slow_time > fast_time
//  [python]
def predict_param_shifts(self, module_kv_map) -> Dict[str, torch.FloatTensor]:
        
        param_shifts = {}
        for module_idx, module_name in enumerate(self.config.inner_params):

            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.device)
            keys = torch.cat([
                module_kv_map[module_idx][idx]["keys"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            values_grad = torch.cat([
                module_kv_map[module_idx][idx]["values_grad"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            value_diffs = torch.empty((0, net.value_size), device = self.config.device)
            for start_idx in range(0, keys.shape[0], self.config.editor_batch_size):
                end_idx = start_idx + self.config.editor_batch_size
                with torch.no_grad():
                    pesudo_keys, pesudo_values_grad = net(
                        keys[start_idx:end_idx],
                        values_grad[start_idx:end_idx],
                        layer_idx
                    )
                    coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diffs = torch.cat((value_diffs, coeffs * pesudo_values_grad))
            with torch.no_grad():
                mat = keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.device)
            param_shift = torch.solve(keys.T @ value_diffs, mat)[0]
            param_shifts[module_name] = param_shift.to(next(self.model.parameters()).device)

        return param_shifts
//  [python]
def predict_param_shifts(self, module_kv_map) -> Dict[str, torch.FloatTensor]:
        
        param_shifts = {}
        for module_idx, module_name in enumerate(self.config.inner_params):

            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.device)
            keys = torch.cat([
                module_kv_map[module_idx][idx]["keys"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            values_grad = torch.cat([
                module_kv_map[module_idx][idx]["values_grad"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            value_diffs = torch.empty((0, net.value_size), device = self.config.device)
            for start_idx in range(0, keys.shape[0], self.config.editor_batch_size):
                end_idx = start_idx + self.config.editor_batch_size
                with torch.no_grad():
                    pesudo_keys, pesudo_values_grad = net(
                        keys[start_idx:end_idx],
                        values_grad[start_idx:end_idx],
                        layer_idx
                    )
                    coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diffs = torch.cat((value_diffs, coeffs * pesudo_values_grad))
            with torch.no_grad():
                mat = keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.device)
            param_shift = torch.linalg.solve(mat, keys.T @ value_diffs)
            param_shifts[module_name] = param_shift.to(next(self.model.parameters()).device)

        return param_shifts
//  [python]
def forward(ctx, A, B, Q, R, T):
        K = [None] * T
        V = [None] * (T+1)

        # for symmetry for gradcheck
        # Q_sym = 0.5*(Q+Q.transpose(-2, -1))
        # R_sym = 0.5*(R+R.transpose(-2, -1))

        A_trans = A.transpose(-2,-1)
        B_trans = B.transpose(-2,-1)
        #initialization for backpropagation
        V[-1] = Q
        A_BK = None

        for i in reversed(range(T)):
            # using torch.solve(B, A) to obtain the solution of A X = B to avoid direct inverse, note it also returns LU
            # print(B_trans, V[i+1], B)
            V_uu_inv_B_trans, _ = torch.solve(B_trans,
                B_trans@V[i+1]@B + R)

            K[i] = V_uu_inv_B_trans@V[i+1]@A
                   
            #riccati difference equation
            # A-BK
            A_BK = A - B@K[i]
            # V = A^T V (A-BK) + Q = A^T V A - A^T V B (B^T V B + R)^{-1} B^T V A + Q
            V[i] = A_trans@V[i+1]@A_BK + Q
        
        ctx.save_for_backward(V[0], K[0], A_BK)
        ctx.T = T
        return V[0]
//  [python]
def forward(ctx, A, B, Q, R, T):
        K = [None] * T
        V = [None] * (T+1)

        # for symmetry for gradcheck
        # Q_sym = 0.5*(Q+Q.transpose(-2, -1))
        # R_sym = 0.5*(R+R.transpose(-2, -1))

        A_trans = A.transpose(-2,-1)
        B_trans = B.transpose(-2,-1)
        #initialization for backpropagation
        V[-1] = Q
        A_BK = None

        for i in reversed(range(T)):
            # using torch.solve(B, A) to obtain the solution of A X = B to avoid direct inverse, note it also returns LU
            # print(B_trans, V[i+1], B)
            # V_uu_inv_B_trans, _ = torch.solve(B_trans,
            #     B_trans@V[i+1]@B + R)
            V_uu_inv_B_trans = torch.linalg.solve(B_trans@V[i+1]@B + R,
                B_trans)

            K[i] = V_uu_inv_B_trans@V[i+1]@A
                   
            #riccati difference equation
            # A-BK
            A_BK = A - B@K[i]
            # V = A^T V (A-BK) + Q = A^T V A - A^T V B (B^T V B + R)^{-1} B^T V A + Q
            V[i] = A_trans@V[i+1]@A_BK + Q
        
        ctx.save_for_backward(V[0], K[0], A_BK)
        ctx.T = T
        return V[0]
//  [python]
def update_hypernet(self, param_shifts: Dict[str, torch.FloatTensor], module_kv_map):
        
        self.opt.zero_grad()
        for module_idx, module_name in enumerate(self.config.inner_params):
            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.device)
            keys = torch.cat([
                module_kv_map[module_idx][idx]["keys"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            values_grad = torch.cat([
                module_kv_map[module_idx][idx]["values_grad"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            module = get_module(self.model, module_name)
            module_grad = module.weight.grad.to(torch.float32).to(self.config.device)
            param_shift = param_shifts[module_name].to(self.config.device)
            if isinstance(module, nn.Linear):
                module_grad = module_grad.T
            with torch.no_grad():
                mat = torch.solve(module_grad, keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.device))[0]
                lamda_grad = - net.lamda(layer_idx).exp() * (mat * param_shift).sum()
            value_diffs_grad = keys @ mat
            (lamda_grad * net.lamda(layer_idx)).backward()
            for start_idx in range(0, keys.shape[0], self.config.editor_batch_size):
                end_idx = start_idx + self.config.editor_batch_size
                pesudo_keys, pesudo_values_grad = net(
                    keys[start_idx:end_idx],
                    values_grad[start_idx:end_idx],
                    layer_idx
                )
                coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diff = coeffs * pesudo_values_grad
                (value_diffs_grad[start_idx:end_idx] * value_diff).sum().backward()
            
        clip_grad_norm_(
            self.net.parameters(),
            self.config.max_grad_norm
        )
        self.opt.step()
//  [python]
def update_hypernet(self, param_shifts: Dict[str, torch.FloatTensor], module_kv_map):
        
        self.opt.zero_grad()
        for module_idx, module_name in enumerate(self.config.inner_params):
            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.device)
            keys = torch.cat([
                module_kv_map[module_idx][idx]["keys"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            values_grad = torch.cat([
                module_kv_map[module_idx][idx]["values_grad"]
                for idx in range(len(module_kv_map[module_idx]))
            ])
            module = get_module(self.model, module_name)
            module_grad = module.weight.grad.to(torch.float32).to(self.config.device)
            param_shift = param_shifts[module_name].to(self.config.device)
            if isinstance(module, nn.Linear):
                module_grad = module_grad.T
            with torch.no_grad():
                mat = torch.linalg.solve(keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.device), module_grad)
                lamda_grad = - net.lamda(layer_idx).exp() * (mat * param_shift).sum()
            value_diffs_grad = keys @ mat
            (lamda_grad * net.lamda(layer_idx)).backward()
            for start_idx in range(0, keys.shape[0], self.config.editor_batch_size):
                end_idx = start_idx + self.config.editor_batch_size
                pesudo_keys, pesudo_values_grad = net(
                    keys[start_idx:end_idx],
                    values_grad[start_idx:end_idx],
                    layer_idx
                )
                coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diff = coeffs * pesudo_values_grad
                (value_diffs_grad[start_idx:end_idx] * value_diff).sum().backward()
            
        clip_grad_norm_(
            self.net.parameters(),
            self.config.max_grad_norm
        )
        self.opt.step()
//  [python]
def fit_poly(self, kernel_func=None, target='hypo', fkine=None, reg=0):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.rbf_kernel = kernel.MultiQuadratic(1) if kernel_func is None else kernel_func
        kmat = self.rbf_kernel(X, X)

        # min_d, min_i = (kmat+torch.eye(len(kmat)).fill_diagonal_(float('inf'))).min(dim=0)
        # plt.plot(range(len(kmat)), min_d)
        # plt.show()
        for c in range(self.num_class):
            c_nonzero = torch.nonzero(self.gains[:, c])
            c_iszero = torch.nonzero(self.gains[:, c] == 0).reshape(-1)
            assert len(c_nonzero) + len(c_iszero) == len(self.gains)
            ridx = c_nonzero.repeat([1, len(c_iszero)]).reshape(-1)
            cidx = c_iszero.repeat([len(c_nonzero), 1]).reshape(-1)
            kmat[ridx, cidx] = 0
            kmat[cidx, ridx] = 0
        self.rbf_nodes = torch.solve(y, kmat+reg*torch.eye(len(kmat), dtype=kmat.dtype))[0]
        for c in range(self.num_class):
            self.rbf_nodes[self.gains == 0] = 0
        assert self.rbf_nodes.shape == (len(self.support_points), self.num_class)
//  [python]
def fit_poly(self, kernel_func=None, target='hypo', fkine=None, reg=0):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.rbf_kernel = kernel.MultiQuadratic(1) if kernel_func is None else kernel_func
        kmat = self.rbf_kernel(X, X)

        # min_d, min_i = (kmat+torch.eye(len(kmat)).fill_diagonal_(float('inf'))).min(dim=0)
        # plt.plot(range(len(kmat)), min_d)
        # plt.show()
        for c in range(self.num_class):
            c_nonzero = torch.nonzero(self.gains[:, c])
            c_iszero = torch.nonzero(self.gains[:, c] == 0).reshape(-1)
            assert len(c_nonzero) + len(c_iszero) == len(self.gains)
            ridx = c_nonzero.repeat([1, len(c_iszero)]).reshape(-1)
            cidx = c_iszero.repeat([len(c_nonzero), 1]).reshape(-1)
            kmat[ridx, cidx] = 0
            kmat[cidx, ridx] = 0
        self.rbf_nodes = torch.linalg.solve(kmat+reg*torch.eye(len(kmat), dtype=kmat.dtype), y)
        for c in range(self.num_class):
            self.rbf_nodes[self.gains == 0] = 0
        assert self.rbf_nodes.shape == (len(self.support_points), self.num_class)
//  [python]
def solve(matrix, vec):
    return torch.solve(vec, matrix)[0]
//  [python]
def solve(matrix, vec):
    return torch.linalg.solve(matrix, vec)
//  [python]
def se3_log_map(transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4, backend=None):
    if backend is None:
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        loguru.logger.warning("!!!!se3_log_map backend is None!!!!")
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        backend = 'pytorch3d'
    if backend == 'pytorch3d':
        return pytorch3d.transforms.se3.se3_log_map(transform, eps, cos_bound)
    elif backend == 'opencv':
        from pytorch3d.transforms.se3 import _se3_V_matrix, _get_se3_V_input
        # from pytorch3d.common.compat import solve
        log_rotation = []
        for tsfm in transform:
            cv2_rot = -cv2.Rodrigues(to_array(tsfm[:3, :3]))[0]
            log_rotation.append(torch.from_numpy(cv2_rot.reshape(-1)).to(transform.device).float())
        log_rotation = torch.stack(log_rotation, dim=0)
        T = transform[:, 3, :3]
        V = _se3_V_matrix(*_get_se3_V_input(log_rotation), eps=eps)
        # log_translation = solve(V, T[:, :, None])[:, :, 0]
        log_translation = torch.solve(T[:, :, None], V)[0][:, :, 0]
        return torch.cat((log_translation, log_rotation), dim=1)
    else:
        raise NotImplementedError()
//  [python]
def se3_log_map(transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4, backend=None):
    if backend is None:
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        loguru.logger.warning("!!!!se3_log_map backend is None!!!!")
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        backend = 'pytorch3d'
    if backend == 'pytorch3d':
        return pytorch3d.transforms.se3.se3_log_map(transform, eps, cos_bound)
    elif backend == 'opencv':
        from pytorch3d.transforms.se3 import _se3_V_matrix, _get_se3_V_input
        # from pytorch3d.common.compat import solve
        log_rotation = []
        for tsfm in transform:
            cv2_rot = -cv2.Rodrigues(to_array(tsfm[:3, :3]))[0]
            log_rotation.append(torch.from_numpy(cv2_rot.reshape(-1)).to(transform.device).float())
        log_rotation = torch.stack(log_rotation, dim=0)
        T = transform[:, 3, :3]
        V = _se3_V_matrix(*_get_se3_V_input(log_rotation), eps=eps)
        # log_translation = solve(V, T[:, :, None])[:, :, 0]
        log_translation = torch.linalg.solve(V, T[:, :, None])[:, :, 0]
        return torch.cat((log_translation, log_rotation), dim=1)
    else:
        raise NotImplementedError()
//  [python]
def bsolve(A, B):
    """Compute solve(A, B) in batch mode, where the first dimension of A can be singleton."""
    if A.dim() == 3 and B.dim() == 2 and A.shape[0] == 1:
        return torch.solve(A.squeeze(0), B.t())[0].t()
    else:
        return torch.solve(A, B)[0]
//  [python]
def bsolve(A, B):
    """Compute solve(A, B) in batch mode, where the first dimension of A can be singleton."""
    if A.dim() == 3 and B.dim() == 2 and A.shape[0] == 1:
        return torch.linalg.solve(A.squeeze(0), B.t()).t()
    else:
        return torch.linalg.solve(A, B)
//  [python]
def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
        return torch.solve(
            self.hessian.to(self.model_device)
            + self.hessian_regularization
            * torch.eye(self.n_parameters, device=self.model_device),
            rhs.T.to(self.model_device),
        )[0].T
//  [python]
def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
        return torch.linalg.solve(
            self.hessian.to(self.model_device)
            + self.hessian_regularization
            * torch.eye(self.n_parameters, device=self.model_device),
            rhs.T.to(self.model_device),
        ).T
//  [python]
def _beta(ZDZ: torch.tensor, ZDy: torch.tensor) -> torch.tensor:
        """
        compute effect size beta = ((U'Z)'D^-1(U'Z))^-1(U'Z)'D^-1(U'y)

        :param ZDZ: precomputed matrix product of (U'Z)'D^-1(U'Z)
        :param ZDy: precomputed matrix product of (U'Z)'D^-1(U'y)

        :return: beta
        """
        return torch.solve(ZDy, ZDZ)[0]
//  [python]
def _beta(ZDZ: torch.tensor, ZDy: torch.tensor) -> torch.tensor:
        """
        compute effect size beta = ((U'Z)'D^-1(U'Z))^-1(U'Z)'D^-1(U'y)

        :param ZDZ: precomputed matrix product of (U'Z)'D^-1(U'Z)
        :param ZDy: precomputed matrix product of (U'Z)'D^-1(U'y)

        :return: beta
        """
        return torch.linalg.solve(ZDZ, ZDy)
//  [python]
def compute_mvdr_bf2(
    covmat_target: torch.Tensor,
    covmat_noise: torch.Tensor,
    ref_mic: Optional[int] = 0,
    eps: Optional[float] = 1e-5,
) -> torch.Tensor:

    num = torch.solve(covmat_target, covmat_noise)[0]
    denom = torch.sum(torch.diagonal(num, dim1=-2, dim2=-1), dim=-1)
    w = num / denom[..., None, None]

    return w[..., ref_mic]
//  [python]
def compute_mvdr_bf2(
    covmat_target: torch.Tensor,
    covmat_noise: torch.Tensor,
    ref_mic: Optional[int] = 0,
    eps: Optional[float] = 1e-5,
) -> torch.Tensor:

    num = torch.linalg.solve(covmat_noise, covmat_target)
    denom = torch.sum(torch.diagonal(num, dim1=-2, dim2=-1), dim=-1)
    w = num / denom[..., None, None]

    return w[..., ref_mic]
//  [python]
def forward(ctx, H, h01, S, s01, ee, method='Lopez-Sancho'):
        # '''
        # gs = [A_l - A_{l,l-1} gs A_{l-1,l}]^{-1}
        # 
        # 1. ee can be a list, to handle a batch of samples
        # '''

        if method == 'GEP':
            gs = calcg0(ee, H, S, h01, s01)
        else:
            h10 = h01.conj().T
            s10 = s01.conj().T
            alpha, beta = h10 - ee * s10, h01 - ee * s01
            eps, epss = H.clone(), H.clone()
            
            converged = False
            iteration = 0
            while not converged:
                iteration += 1
                oldeps, oldepss = eps.clone(), epss.clone()
                oldalpha, oldbeta = alpha.clone(), beta.clone()
                tmpa = torch.solve(oldalpha, ee * S - oldeps)[0]
                tmpb = torch.solve(oldbeta, ee * S - oldeps)[0]

                alpha, beta = torch.mm(oldalpha, tmpa), torch.mm(oldbeta, tmpb)
                eps = oldeps + torch.mm(oldalpha, tmpb) + torch.mm(oldbeta, tmpa)

                epss = oldepss + torch.mm(oldbeta, tmpa)
                LopezConvTest = torch.max(alpha.abs() + beta.abs())

                if iteration == 101:
                    log.error("Lopez-scheme not converged after 100 iteration.")

                if LopezConvTest < 1.0e-40:
                    gs = torch.inverse(ee * S - epss)

                    test = ee * S - H - torch.mm(ee * s01 - h01, gs.mm(ee * s10 - h10))
                    myConvTest = torch.max((test.mm(gs) - torch.eye(H.shape[0], dtype=h01.dtype)).abs())
                    if myConvTest < 1.0e-6:
                        converged = True
                        if myConvTest > 1.0e-8:
                            log.warning("Lopez-scheme not-so-well converged at E = %.4f eV:" % ee.real.item() + str(myConvTest.item()))
                    else:
                        log.error("Lopez-Sancho %.8f " % myConvTest.item() +
                              "Error: gs iteration {0}".format(iteration))
                        raise ArithmeticError("Criteria not met. Please check output...")
                    
        ctx.save_for_backward(gs, H, h01, S, s01, ee)
        return gs
//  [python]
def forward(ctx, H, h01, S, s01, ee, method='Lopez-Sancho'):
        # '''
        # gs = [A_l - A_{l,l-1} gs A_{l-1,l}]^{-1}
        # 
        # 1. ee can be a list, to handle a batch of samples
        # '''

        if method == 'GEP':
            gs = calcg0(ee, H, S, h01, s01)
        else:
            h10 = h01.conj().T
            s10 = s01.conj().T
            alpha, beta = h10 - ee * s10, h01 - ee * s01
            eps, epss = H.clone(), H.clone()
            
            converged = False
            iteration = 0
            while not converged:
                iteration += 1
                oldeps, oldepss = eps.clone(), epss.clone()
                oldalpha, oldbeta = alpha.clone(), beta.clone()
                tmpa = tLA.solve(ee * S - oldeps, oldalpha)
                tmpb = tLA.solve(ee * S - oldeps, oldbeta)

                alpha, beta = torch.mm(oldalpha, tmpa), torch.mm(oldbeta, tmpb)
                eps = oldeps + torch.mm(oldalpha, tmpb) + torch.mm(oldbeta, tmpa)

                epss = oldepss + torch.mm(oldbeta, tmpa)
                LopezConvTest = torch.max(alpha.abs() + beta.abs())

                if iteration == 101:
                    log.error("Lopez-scheme not converged after 100 iteration.")

                if LopezConvTest < 1.0e-40:
                    gs = (ee * S - epss).inverse()

                    test = ee * S - H - torch.mm(ee * s01 - h01, gs.mm(ee * s10 - h10))
                    myConvTest = torch.max((test.mm(gs) - torch.eye(H.shape[0], dtype=h01.dtype)).abs())
                    if myConvTest < 1.0e-6:
                        converged = True
                        if myConvTest > 1.0e-8:
                            log.warning("Lopez-scheme not-so-well converged at E = %.4f eV:" % ee.real.item() + str(myConvTest.item()))
                    else:
                        log.error("Lopez-Sancho %.8f " % myConvTest.item() +
                              "Error: gs iteration {0}".format(iteration))
                        raise ArithmeticError("Criteria not met. Please check output...")
                    
        ctx.save_for_backward(gs, H, h01, S, s01, ee)
        return gs
//  [python]
def linalg_solve(A, B):
    return torch.solve(A.float(), B.float)[0].to(A.dtype)
//  [python]
def linalg_solve(A, B):
    return torch.linalg.solve(A.float(), B.float()).to(A.dtype)
//  [python]
def compute_gev_bf(
    covmat_target: torch.Tensor,
    covmat_noise: torch.Tensor,
    ref_mic: Optional[int] = 0,
):

    # now compute the GEV beamformers
    # eigenvalues are listed in ascending order
    eigval, eigvec = eigh(covmat_target, covmat_noise)

    # compute the scale
    rhs = torch.eye(eigvec.shape[-1], dtype=eigvec.dtype, device=eigvec.device)
    rhs = rhs[:, -1, None]
    steering_vector = torch.solve(rhs, hermite(eigvec))[0]
    scale = steering_vector[..., ref_mic, :].transpose(-3, -2)

    # (batch, freq, sources, channels)
    gev_bf = torch.conj(eigvec[..., -1].transpose(-3, -2)) * scale

    return gev_bf
//  [python]
def compute_gev_bf(
    covmat_target: torch.Tensor,
    covmat_noise: torch.Tensor,
    ref_mic: Optional[int] = 0,
):

    # now compute the GEV beamformers
    # eigenvalues are listed in ascending order
    eigval, eigvec = eigh(covmat_target, covmat_noise)

    # compute the scale
    rhs = torch.eye(eigvec.shape[-1], dtype=eigvec.dtype, device=eigvec.device)
    rhs = rhs[:, -1, None]
    steering_vector = torch.linalg.solve(hermite(eigvec), rhs)
    scale = steering_vector[..., ref_mic, :].transpose(-3, -2)

    # (batch, freq, sources, channels)
    gev_bf = torch.conj(eigvec[..., -1].transpose(-3, -2)) * scale

    return gev_bf
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.solve(
            B.unsqueeze(-1), A_backwards
        )[0].squeeze(-1)  # (... N)

    dA = torch.solve(A_forwards, A_backwards)[0]  # (... N N)
    return dA, dB
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.linalg.solve(
            A_backwards, B.unsqueeze(-1)
        ).squeeze(-1)  # (... N)

    dA = torch.linalg.solve(A_backwards, A_forwards)  # (... N N)
    return dA, dB
//  [python]
def forward(self, x):
        eps = 8.0
        outputs = self.model(x)
        softmax_prob = F.softmax(outputs, dim=1)
        smax_inp = softmax_prob

        eye = torch.eye(1000).to(outputs.device)
        eye = eye.reshape((1, 1000, 1000))
        eye = eye.repeat(outputs.shape[0], 1, 1)
        t2 = eps * torch.diag_embed(smax_inp)
        smax_inp = torch.unsqueeze(smax_inp, 2)
        t3 = eps * torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
        matrix = eye + t2 - t3
        y_star = torch.solve(smax_inp, matrix)[0]
        y_star = torch.squeeze(y_star)

        pseudo_prob = y_star
        loss = torch.logsumexp(outputs, dim=1) - (
            pseudo_prob * outputs - eps * pseudo_prob * (1 - softmax_prob)
        ).sum(dim=1)
        loss = loss.mean()
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad(set_to_none=True)
        return outputs
//  [python]
def forward(self, x):
        eps = 8.0
        outputs = self.model(x)
        softmax_prob = F.softmax(outputs, dim=1)
        smax_inp = softmax_prob

        eye = torch.eye(1000).to(outputs.device)
        eye = eye.reshape((1, 1000, 1000))
        eye = eye.repeat(outputs.shape[0], 1, 1)
        t2 = eps * torch.diag_embed(smax_inp)
        smax_inp = torch.unsqueeze(smax_inp, 2)
        t3 = eps * torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
        matrix = eye + t2 - t3
        y_star = torch.linalg.solve(matrix, smax_inp)
        y_star = torch.squeeze(y_star)

        pseudo_prob = y_star
        loss = torch.logsumexp(outputs, dim=1) - (
            pseudo_prob * outputs - eps * pseudo_prob * (1 - softmax_prob)
        ).sum(dim=1)
        loss = loss.mean()
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad(set_to_none=True)
        return outputs
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj() if self.Q is None else _r2c(self.Q)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2 *
             contract('r h n, h n, s h n -> h r s', Q, D, P).real)  # (H r r)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        R = torch.solve(Q_D, R.to(Q_D))[0]  # (H r N)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (r H N)
            "P": P,  # (r H N)
            "Q": Q,  # (r H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj() if self.Q is None else _r2c(self.Q)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2 *
             contract('r h n, h n, s h n -> h r s', Q, D, P).real)  # (H r r)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        R = torch.linalg.solve(R.to(Q_D), Q_D)  # (H r N)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (r H N)
            "P": P,  # (r H N)
            "Q": Q,  # (r H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C))  # (H C N)

        dC = torch.solve(
            C.unsqueeze(-1),
            I - dA_L.transpose(-1, -2),
        ).solution.squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(
                    V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def setup_step(self, mode='dense'):
        """ Set up dA, dB, dC discretized parameters for stepping """
        self._setup_state()

        # Calculate original C
        dA_L = power(self.L, self.dA)
        I = torch.eye(self.dA.size(-1)).to(dA_L)
        C = _conj(_r2c(self.C))  # (H C N)

        dC = torch.linalg.solve(
            I - dA_L.transpose(-1, -2),
            C.unsqueeze(-1),
        ).squeeze(-1)
        self.dC = dC

        # Do special preprocessing for different step modes

        self._step_mode = mode
        if mode == 'linear':
            # Linear case: special step function for the state, we need to handle output
            # use conjugate symmetry by default, which affects the output projection
            self.dC = 2*self.dC[:, :, :self.N]
        elif mode == 'diagonal':
            # Eigendecomposition of the A matrix
            L, V = torch.linalg.eig(self.dA)
            V_inv = torch.linalg.inv(V)
            # Check that the eigendedecomposition is correct
            if self.verbose:
                print("Diagonalization error:", torch.dist(
                    V @ torch.diag_embed(L) @ V_inv, self.dA))

            # Change the parameterization to diagonalize
            self.dA = L
            self.dB = contract('h n m, h m -> h n', V_inv, self.dB)
            self.dC = contract('h n m, c h n -> c h m', V, self.dC)

        elif mode == 'dense':
            pass
        else:
            raise NotImplementedError(
                "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}")
//  [python]
def get_correction_factor(self, double=False):

        # Get the parameters which are transformed (do not use noisy training on these params)
        d = self.get_diagonal()  # (Q, N)
        p = self.get_lowrank()  # (Q, N)
        t = self.get_timestep()  # (Q, H)

        identity = torch.eye(self.hidden_dim).to(d).unsqueeze(0).unsqueeze(0)

        # Get continous matrix (H, N, N)
        A = 0.50 * contract("qh,qnm->qhnm", t, self.get_transition(d, p))

        # Get discretized A naively
        # print("Solving dA = solve(identity - A, identity + A)", A.size(), A.type(), A.device)
        dA = torch.solve(identity + A, identity - A)[0]

        # Correction factor
        # Get identity (1, N, N)
        if double:
            return identity + torch.matrix_power(dA, self.maxlen)
        return identity - torch.matrix_power(dA, self.maxlen)
//  [python]
def get_correction_factor(self, double=False):

        # Get the parameters which are transformed (do not use noisy training on these params)
        d = self.get_diagonal()  # (Q, N)
        p = self.get_lowrank()  # (Q, N)
        t = self.get_timestep()  # (Q, H)

        identity = torch.eye(self.hidden_dim).to(d).unsqueeze(0).unsqueeze(0)

        # Get continous matrix (H, N, N)
        A = 0.50 * contract("qh,qnm->qhnm", t, self.get_transition(d, p))

        # Get discretized A naively
        # print("Solving dA = solve(identity - A, identity + A)", A.size(), A.type(), A.device)
        dA = torch.linalg.solve(identity - A, identity + A)

        # Correction factor
        # Get identity (1, N, N)
        if double:
            return identity + torch.matrix_power(dA, self.maxlen)
        return identity - torch.matrix_power(dA, self.maxlen)
//  [python]
def get_geometry(
            self,
            fH,
            fW,
            lidar2img,
            img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        if self.frustum == None:
            self.frustum = self.create_frustum(fH, fW, img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]

        points = self.frustum.view(1, 1, self.D, fH, fW, 3) \
            .repeat(B, N, 1, 1, 1, 1)
        lidar2img = lidar2img.view(B, N, 1, 1, 1, 4, 4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.solve(points.unsqueeze(-1).to(torch.float32),
                                    lidar2img.to(torch.float32))[0].squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points