jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def predict(self, x):
        """Computes the predicted means and variances at the given location.

        The posterior means and variances are returned if `fit()` has already been
        called otherwise the prior means and variances are returned.

        Args:
            x: The location to predict at.
        Returns
            The predicted means and variances.
        """
        if self.train_x is None:
            return self.mean(x), self.kernel(x, x)

        # Compute means and covariances.
        train_mu, x_mu = self.mean(self.train_x), self.mean(x)
        train_sig = self.kernel(
            self.train_x, self.train_x
        ) + self.noise_var * torch.eye(self.train_x.shape[0])
        x_sig, cross_sig = self.kernel(x, x), self.kernel(self.train_x, x)

        # Compute posterior mean and covariance.
        solved = torch.linalg.solve(train_sig, cross_sig).T
        mu = x_mu + solved @ (self.train_y - train_mu)
        sig = x_sig - (solved @ cross_sig)

        return mu, sig
//  [python]
def get_mwf_vector(
    psd_s,
    psd_n,
    reference_vector: Union[torch.Tensor, int],
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
):
    """Return the MWF (Minimum Multi-channel Wiener Filter) vector:

        h = (Npsd^-1 @ Spsd) @ u

    Args:
        psd_s (torch.complex64):
            speech covariance matrix (..., F, C, C)
        psd_n (torch.complex64):
            power-normalized observation covariance matrix (..., F, C, C)
        reference_vector (torch.Tensor or int): (..., C) or scalar
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: D400
    if diagonal_loading:
        psd_n = tik_reg(psd_n, reg=diag_eps, eps=eps)

    ws = torch.solve(psd_s, psd_n)[0]
    # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
    if isinstance(reference_vector, int):
        beamform_vector = ws[..., reference_vector]
    else:
        beamform_vector = torch.einsum(
            "...fec,...c->...fe", ws, reference_vector.to(dtype=ws.dtype)
        )
    return beamform_vector
//  [python]
def get_mwf_vector(
    psd_s,
    psd_n,
    reference_vector: Union[torch.Tensor, int],
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
):
    """Return the MWF (Minimum Multi-channel Wiener Filter) vector:

        h = (Npsd^-1 @ Spsd) @ u

    Args:
        psd_s (torch.complex64):
            speech covariance matrix (..., F, C, C)
        psd_n (torch.complex64):
            power-normalized observation covariance matrix (..., F, C, C)
        reference_vector (torch.Tensor or int): (..., C) or scalar
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: D400
    if diagonal_loading:
        psd_n = tik_reg(psd_n, reg=diag_eps, eps=eps)

    ws = torch.linalg.solve(psd_n, psd_s)
    # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
    if isinstance(reference_vector, int):
        beamform_vector = ws[..., reference_vector]
    else:
        beamform_vector = torch.einsum(
            "...fec,...c->...fe", ws, reference_vector.to(dtype=ws.dtype)
        )
    return beamform_vector
//  [python]
def fit_window_point(self, x, y, history=10, type='slope'):

        nonzero_index = torch.squeeze(torch.nonzero(y), dim=1)
        y = torch.index_select(y, 0, nonzero_index)
        x = torch.index_select(x, 0, nonzero_index)

        ws = torch.flip(1 ** torch.arange(len(y)), dims=[0])
        ws = ws.float()

        if len(y) >= 2:
            if type == 'slope':
                X = torch.stack((x, torch.ones_like(x))).T
                X = X.float()
            else:
                X = torch.stack((x ** 2, x, torch.ones_like(x))).T
            w = torch.solve(X.T @ (ws[:, None] * X), X.T @ (ws * y))[0]

            result = w[0]
        else:
            result = 0.0

        return result
//  [python]
def fit_window_point(self, x, y, history=10, type='slope'):

        nonzero_index = torch.squeeze(torch.nonzero(y), dim=1)
        y = torch.index_select(y, 0, nonzero_index)
        x = torch.index_select(x, 0, nonzero_index)

        ws = torch.flip(1 ** torch.arange(len(y)), dims=[0])
        ws = ws.float()

        if len(y) >= 2:
            if type == 'slope':
                X = torch.stack((x, torch.ones_like(x))).T
                X = X.float()
            else:
                X = torch.stack((x ** 2, x, torch.ones_like(x))).T
            w = torch.linalg.solve(X.T @ (ws[:, None] * X), X.T @ (ws * y))

            result = w[0]
        else:
            result = 0.0

        return result
//  [python]
def torch_solve_qp_backwards(dl_dz, sol_mats, n_eq, n_ineq):
    # --- prep:
    n_batch = dl_dz.shape[0]
    n_x = dl_dz.shape[1]
    n_con = n_eq + n_ineq
    zeros = torch.zeros(n_batch, n_con, 1)

    # --- rhs:
    rhs = torch.cat((-dl_dz, zeros), dim=1)
    back_sol = torch.solve(rhs, sol_mats)[0]

    # --- unpack solution:
    dx = back_sol[:, :n_x, :]
    dlam = None
    dnu = None
    if n_ineq > 0:
        dlam = back_sol[:, n_x:(n_x + n_ineq), :]
    if n_eq > 0:
        dnu = back_sol[:, (n_x + n_ineq):(n_x + n_ineq + n_eq), :]

    diff_list = {"dx": dx, "dlam": dlam, "dnu": dnu}
    return diff_list
//  [python]
def torch_solve_qp_backwards(dl_dz, sol_mats, n_eq, n_ineq):
    # --- prep:
    n_batch = dl_dz.shape[0]
    n_x = dl_dz.shape[1]
    n_con = n_eq + n_ineq
    zeros = torch.zeros(n_batch, n_con, 1)

    # --- rhs:
    rhs = torch.cat((-dl_dz, zeros), dim=1)
    back_sol = torch.linalg.solve(sol_mats, rhs)

    # --- unpack solution:
    dx = back_sol[:, :n_x, :]
    dlam = None
    dnu = None
    if n_ineq > 0:
        dlam = back_sol[:, n_x:(n_x + n_ineq), :]
    if n_eq > 0:
        dnu = back_sol[:, (n_x + n_ineq):(n_x + n_ineq + n_eq), :]

    diff_list = {"dx": dx, "dlam": dlam, "dnu": dnu}
    return diff_list
//  [python]
def __init__(self, d: int, tau: Union[Tensor, float], w_init: Tensor = None):
        super().__init__()
        # fixing gamma parameterization mixup
        tau = 1 / math.sqrt(tau)
        if isinstance(tau, float):
            tau = torch.tensor(tau)
        self.register_buffer("tau", tau)
        tau_data = torch.tensor(
            [1.0, 1 / 1e-1, 1 / 1e-2, 1 / 1e-3, 1 / 1e-4, 1 / 1e-5]
        ).log()
        tau_data = torch.stack([torch.ones(6), tau_data], dim=-1)
        mu_data = torch.tensor(
            [-1.6932, -6.2983, -10.9035, -15.5087, -20.1138, -24.7190]
        )
        # scale_data = (tau_data.log() @ mu_data) / (tau_data.log() @ tau_data.log())
        scale = torch.solve(tau_data.t() @ mu_data, tau_data.t() @ tau_data)[0]
        if w_init is not None:
            assert len(w_init) == d, "w_init must be a vector of length d."
            global_scale_init = -1.6931 * torch.ones(1)
            noise = -5.0
        else:
            # initializing such that kl-divergence is minimized
            global_scale_init = scale[0] + scale[1] * tau.log() * torch.ones(1)
            noise = 0.3466
        # global_scale_init = -1.6931 * torch.ones(1)
        # data fit for tau
        self.s_a = LogNormalMeanFieldVariational(
            global_scale_init, noise + torch.randn(1) * 1e-4
        )
        self.s_b = LogNormalMeanFieldVariational(
            1.6931 * torch.ones(1), noise + torch.randn(1) * 1e-4
        )
        self.gamma_a = LogNormalMeanFieldVariational(
            -1.6931 * torch.ones(d), noise + torch.randn(d) * 1e-4
        )
        self.gamma_b = LogNormalMeanFieldVariational(
            +1.6931 * torch.ones(d), noise + torch.randn(d) * 1e-4
        )
        if w_init is None:
            w_init = torch.zeros(d)
            w_tilde_noise = -0.0
        else:
            w_tilde_noise = -6.0
        self.w_tilde = NormalMeanFieldVariational(
            w_init, w_tilde_noise + torch.randn(d) * 1e-6
        )
        self.register_buffer("sparse_index", torch.ones(d, dtype=torch.bool))
        self.register_buffer("purning_tol", torch.tensor(0.0))
//  [python]
def __init__(self, d: int, tau: Union[Tensor, float], w_init: Tensor = None):
        super().__init__()
        # fixing gamma parameterization mixup
        tau = 1 / math.sqrt(tau)
        if isinstance(tau, float):
            tau = torch.tensor(tau)
        self.register_buffer("tau", tau)
        tau_data = torch.tensor(
            [1.0, 1 / 1e-1, 1 / 1e-2, 1 / 1e-3, 1 / 1e-4, 1 / 1e-5]
        ).log()
        tau_data = torch.stack([torch.ones(6), tau_data], dim=-1)
        mu_data = torch.tensor(
            [-1.6932, -6.2983, -10.9035, -15.5087, -20.1138, -24.7190]
        )
        # scale_data = (tau_data.log() @ mu_data) / (tau_data.log() @ tau_data.log())
        scale = torch.linalg.solve(tau_data.t() @ tau_data, (tau_data.t() @ mu_data))
        if w_init is not None:
            assert len(w_init) == d, "w_init must be a vector of length d."
            global_scale_init = -1.6931 * torch.ones(1)
            noise = -5.0
        else:
            # initializing such that kl-divergence is minimized
            global_scale_init = scale[0] + scale[1] * tau.log() * torch.ones(1)
            noise = 0.3466
        # global_scale_init = -1.6931 * torch.ones(1)
        # data fit for tau
        self.s_a = LogNormalMeanFieldVariational(
            global_scale_init, noise + torch.randn(1) * 1e-4
        )
        self.s_b = LogNormalMeanFieldVariational(
            1.6931 * torch.ones(1), noise + torch.randn(1) * 1e-4
        )
        self.gamma_a = LogNormalMeanFieldVariational(
            -1.6931 * torch.ones(d), noise + torch.randn(d) * 1e-4
        )
        self.gamma_b = LogNormalMeanFieldVariational(
            +1.6931 * torch.ones(d), noise + torch.randn(d) * 1e-4
        )
        if w_init is None:
            w_init = torch.zeros(d)
            w_tilde_noise = -0.0
        else:
            w_tilde_noise = -6.0
        self.w_tilde = NormalMeanFieldVariational(
            w_init, w_tilde_noise + torch.randn(d) * 1e-6
        )
        self.register_buffer("sparse_index", torch.ones(d, dtype=torch.bool))
        self.register_buffer("purning_tol", torch.tensor(0.0))
//  [python]
def get_lcmv_vector_with_rtf(
    psd_n: torch.Tensor,
    rtf_mat: torch.Tensor,
    reference_vector: Union[int, torch.Tensor, None] = None,
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Return the LCMV (Linearly Constrained Minimum Variance) vector
        calculated with RTF:

        h = (Npsd^-1 @ rtf_mat) @ (rtf_mat^H @ Npsd^-1 @ rtf_mat)^-1 @ p

    Reference:
        H. L. Van Trees, “Optimum array processing: Part IV of detection, estimation,
        and modulation theory,” John Wiley & Sons, 2004. (Chapter 6.7)

    Args:
        psd_n (torch.complex64):
            observation/noise covariance matrix (..., F, C, C)
        rtf_mat (torch.complex64):
            RTF matrix (..., F, C, num_spk)
        reference_vector (torch.Tensor or int): (..., num_spk) or scalar
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: H405, D205, D400
    if diagonal_loading:
        psd_n = tik_reg(psd_n, reg=diag_eps, eps=eps)

    # numerator: (..., C_1, C_2) x (..., C_2, num_spk) -> (..., C_1, num_spk)
    numerator = torch.solve(rtf_mat, psd_n)[0]
    denominator = torch.matmul(rtf_mat.conj().transpose(-1, -2), numerator)
    if isinstance(reference_vector, int):
        ws = torch.inverse(denominator)[..., reference_vector, None]
    else:
        ws = torch.solve(reference_vector, denominator)[0]
    beamforming_vector = torch.matmul(numerator, ws).squeeze(-1)
    return beamforming_vector
//  [python]
def get_lcmv_vector_with_rtf(
    psd_n: torch.Tensor,
    rtf_mat: torch.Tensor,
    reference_vector: Union[int, torch.Tensor, None] = None,
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Return the LCMV (Linearly Constrained Minimum Variance) vector
        calculated with RTF:

        h = (Npsd^-1 @ rtf_mat) @ (rtf_mat^H @ Npsd^-1 @ rtf_mat)^-1 @ p

    Reference:
        H. L. Van Trees, “Optimum array processing: Part IV of detection, estimation,
        and modulation theory,” John Wiley & Sons, 2004. (Chapter 6.7)

    Args:
        psd_n (torch.complex64):
            observation/noise covariance matrix (..., F, C, C)
        rtf_mat (torch.complex64):
            RTF matrix (..., F, C, num_spk)
        reference_vector (torch.Tensor or int): (..., num_spk) or scalar
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: H405, D205, D400
    if diagonal_loading:
        psd_n = tik_reg(psd_n, reg=diag_eps, eps=eps)

    # numerator: (..., C_1, C_2) x (..., C_2, num_spk) -> (..., C_1, num_spk)
    numerator = torch.linalg.solve(psd_n, rtf_mat)
    denominator = torch.matmul(rtf_mat.conj().transpose(-1, -2), numerator)
    if isinstance(reference_vector, int):
        ws = denominator.inverse()[..., reference_vector, None]
    else:
        ws = torch.linalg.solve(denominator, reference_vector)
    beamforming_vector = torch.matmul(numerator, ws).squeeze(-1)
    return beamforming_vector
//  [python]
def guide(self, x_ng: torch.Tensor) -> None:
        if self.ppca_flavor == "marginalized":
            return

        with pyro.plate("cells", size=self.n_obs, subsample_size=x_ng.shape[0]):
            V_gk = torch.solve(self.W_kg.T, self.M_kk)[0].T
            D_k = self.sigma / torch.sqrt(torch.diag(self.M_kk))
            pyro.sample("z", dist.Normal((x_ng - self.mean_g) @ V_gk, D_k).to_event(1))
//  [python]
def guide(self, x_ng: torch.Tensor) -> None:
        if self.ppca_flavor == "marginalized":
            return

        with pyro.plate("cells", size=self.n_obs, subsample_size=x_ng.shape[0]):
            V_gk = torch.linalg.solve(self.M_kk, self.W_kg).T
            D_k = self.sigma / torch.sqrt(torch.diag(self.M_kk))
            pyro.sample("z", dist.Normal((x_ng - self.mean_g) @ V_gk, D_k).to_event(1))
//  [python]
def solve(self, f, bc):
        """
        This method is used to solve the PDE with the provided source function and boundary condition.

        Parameters
        ----------
        f : tensor/array/list
            The source function for the PDE.
        bc : tensor/array/list
            The boundary condition for the PDE.

        Returns
        -------
        tuple
            The tuple contains the following elements:

            u_pred : tensor
                The complete solution of the PDE.
            u_pde_pred : tensor
                The solution of the PDE inside the domain.
            u_bc_pred : tensor
                The solution for the boundary condition.
            f_pred : tensor
                The predicted source function.
            runtime : float
                The time it took the method to run in seconds.
        """
        self.f, self.bc = format_input([f, bc], self.precision, self.device)

        w_outs = []
        biases = np.empty(len(self.lambdas_pde))

        t0 = time.perf_counter()

        RHSs = torch.matmul(self.RHSs, self.f)
        self.bct_ones = torch.sum(self.bc).reshape(1, 1)

        for i, l in enumerate(self.lambdas_pde):
            self.w_out = torch.solve(RHSs[i], self.LHSs[i])[0]
            self.bias = - 1 / self.NdO * (self.Ht_bc_ones @ self.w_out - self.bct_ones)

            # Just calculates the loss if multiple lambdas are provided to find the best one.
            if self.n_lambdas > 1:
                u_pred_bc = torch.matmul(self.H_bc, self.w_out) + self.bias
                f_pred = torch.matmul(self.DH, self.w_out)
                self.Ls_pde[i] = torch.mean((f_pred - self.f) ** 2).item()
                self.Ls_bc[i] = torch.mean((u_pred_bc - self.bc) ** 2).item()
                self.Ls[i] = self.Ls_pde[i] + self.Ls_bc[i]

                w_outs.append(self.w_out)
                biases[i] = self.bias

        if self.n_lambdas > 1:
            minimum = np.argmin(self.Ls)
            self.lambda_pde = self.lambdas_pde[minimum]
            self.w_out = w_outs[minimum]
            self.bias = biases[minimum]
        else:
            self.lambda_pde = self.lambdas_pde[0]

        self.u_pde_pred = torch.add(torch.matmul(self.H, self.w_out), self.bias)

        self.f_pred = torch.matmul(self.DH, self.w_out)
        self.u_bc_pred = torch.matmul(self.H_bc, self.w_out) + self.bias
        self.u_pred = torch.cat([self.u_pde_pred, self.u_bc_pred])

        t1 = time.perf_counter()
        if self.verbose > 0:
            print('\nRun Successful:', t1 - t0, 'seconds\n')

        return self.u_pred, self.u_pde_pred, self.u_bc_pred, self.f_pred, t1 - t0
//  [python]
def solve(self, f, bc):
        """
        This method is used to solve the PDE with the provided source function and boundary condition.

        Parameters
        ----------
        f : tensor/array/list
            The source function for the PDE.
        bc : tensor/array/list
            The boundary condition for the PDE.

        Returns
        -------
        tuple
            The tuple contains the following elements:

            u_pred : tensor
                The complete solution of the PDE.
            u_pde_pred : tensor
                The solution of the PDE inside the domain.
            u_bc_pred : tensor
                The solution for the boundary condition.
            f_pred : tensor
                The predicted source function.
            runtime : float
                The time it took the method to run in seconds.
        """
        self.f, self.bc = format_input([f, bc], self.precision, self.device)

        w_outs = []
        biases = np.empty(len(self.lambdas_pde))

        t0 = time.perf_counter()

        RHSs = torch.matmul(self.RHSs, self.f)
        self.bct_ones = torch.sum(self.bc).reshape(1, 1)

        for i, l in enumerate(self.lambdas_pde):
            self.w_out = torch.linalg.solve(self.LHSs[i], RHSs[i])
            self.bias = - 1 / self.NdO * (self.Ht_bc_ones @ self.w_out - self.bct_ones)

            # Just calculates the loss if multiple lambdas are provided to find the best one.
            if self.n_lambdas > 1:
                u_pred_bc = torch.matmul(self.H_bc, self.w_out) + self.bias
                f_pred = torch.matmul(self.DH, self.w_out)
                self.Ls_pde[i] = torch.mean((f_pred - self.f) ** 2).item()
                self.Ls_bc[i] = torch.mean((u_pred_bc - self.bc) ** 2).item()
                self.Ls[i] = self.Ls_pde[i] + self.Ls_bc[i]

                w_outs.append(self.w_out)
                biases[i] = self.bias

        if self.n_lambdas > 1:
            minimum = np.argmin(self.Ls)
            self.lambda_pde = self.lambdas_pde[minimum]
            self.w_out = w_outs[minimum]
            self.bias = biases[minimum]
        else:
            self.lambda_pde = self.lambdas_pde[0]

        self.u_pde_pred = torch.add(torch.matmul(self.H, self.w_out), self.bias)

        self.f_pred = torch.matmul(self.DH, self.w_out)
        self.u_bc_pred = torch.matmul(self.H_bc, self.w_out) + self.bias
        self.u_pred = torch.cat([self.u_pde_pred, self.u_bc_pred])

        t1 = time.perf_counter()
        if self.verbose > 0:
            print('\nRun Successful:', t1 - t0, 'seconds\n')

        return self.u_pred, self.u_pde_pred, self.u_bc_pred, self.f_pred, t1 - t0
//  [python]
def optimize(A, B):
    """
    least square optimizer:
        A @ u.T = B
    """
    # L = torch.cholesky(A + torch.eye(A.shape[1]).to(device) * 1e-8)
    # y = torch.triangular_solve(B, L.T, upper=True)[0]
    # u = torch.triangular_solve(y, L, upper=False)[0]
    u, _ = torch.solve(B, A)
    return u
//  [python]
def optimize(A, B):
    """
    least square optimizer:
        A @ u.T = B
    """
    u, _ = torch.linalg.solve(A, B)
    return u
//  [python]
def predict(self, x_ng: torch.Tensor, var_names_g: np.ndarray) -> dict[str, np.ndarray | torch.Tensor]:
        """
        Centering and embedding of the input data ``x_ng`` into the principal component space.

        .. note::
           Gradients are disabled, used for inference only.

        Args:
            x_ng:
                Gene counts matrix.
            var_names_g:
                The list of the variable names in the input data.

        Returns:
            A dictionary with the following keys:

            - ``z_nk``: Embedding of the input data into the principal component space.
        """
        assert_columns_and_array_lengths_equal("x_ng", x_ng, "var_names_g", var_names_g)
        assert_arrays_equal("var_names_g", var_names_g, "var_names_g", self.var_names_g)

        V_gk = torch.solve(self.W_kg, self.M_kk)[0].T
        z_nk = (x_ng - self.mean_g) @ V_gk
        return {"z_nk": z_nk}
//  [python]
def predict(self, x_ng: torch.Tensor, var_names_g: np.ndarray) -> dict[str, np.ndarray | torch.Tensor]:
        """
        Centering and embedding of the input data ``x_ng`` into the principal component space.

        .. note::
           Gradients are disabled, used for inference only.

        Args:
            x_ng:
                Gene counts matrix.
            var_names_g:
                The list of the variable names in the input data.

        Returns:
            A dictionary with the following keys:

            - ``z_nk``: Embedding of the input data into the principal component space.
        """
        assert_columns_and_array_lengths_equal("x_ng", x_ng, "var_names_g", var_names_g)
        assert_arrays_equal("var_names_g", var_names_g, "var_names_g", self.var_names_g)

        V_gk = torch.linalg.solve(self.M_kk, self.W_kg).T
        z_nk = (x_ng - self.mean_g) @ V_gk
        return {"z_nk": z_nk}
//  [python]
def laplacian_solve(L: torch.Tensor, y_hat: torch.Tensor, lambda_param: float = None):
    """Solves a system of linear equations to get the post-processed output.
    The system of linear equation it solves is:
    :math:`\hat{{f}} = {(I + \lambda * L)}^{-1} \hat{y}`

    Parameters
    ------------
        L: torch.Tensor
            Laplacian matrix
        y_hat: torch.Tensor
            Model predicted output class probabilities
        lambda_param: float
            Weight for the laplacian regularizer

    Returns
    ----------
        y: torch.Tensor
            Post-processed solution according to the equation above
    """

    lambda_param = 1.0 if lambda_param is None else lambda_param
    n = L.shape[0]
    y = torch.solve(y_hat, lambda_param * L + torch.eye(n))[0]

    return y
//  [python]
def laplacian_solve(L: torch.Tensor, y_hat: torch.Tensor, lambda_param: float = None):
    """Solves a system of linear equations to get the post-processed output.
    The system of linear equation it solves is:
    :math:`\hat{{f}} = {(I + \lambda * L)}^{-1} \hat{y}`

    Parameters
    ------------
        L: torch.Tensor
            Laplacian matrix
        y_hat: torch.Tensor
            Model predicted output class probabilities
        lambda_param: float
            Weight for the laplacian regularizer

    Returns
    ----------
        y: torch.Tensor
            Post-processed solution according to the equation above
    """

    lambda_param = 1.0 if lambda_param is None else lambda_param
    n = L.shape[0]
    y = torch.linalg.solve(lambda_param * L + torch.eye(n), y_hat)

    return y
//  [python]
def get_gev_vector(
    psd_noise: torch.Tensor,
    psd_speech: torch.Tensor,
    mode="power",
    reference_vector: Union[int, torch.Tensor] = 0,
    iterations: int = 3,
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Return the generalized eigenvalue (GEV) beamformer vector:

        psd_speech @ h = lambda * psd_noise @ h

    Reference:
        Blind acoustic beamforming based on generalized eigenvalue decomposition;
        E. Warsitz and R. Haeb-Umbach, 2007.

    Args:
        psd_noise (torch.complex64):
            noise covariance matrix (..., F, C, C)
        psd_speech (torch.complex64):
            speech covariance matrix (..., F, C, C)
        mode (str): one of ("power", "evd")
            "power": power method
            "evd": eigenvalue decomposition
        reference_vector (torch.Tensor or int): (..., C) or scalar
        iterations (int): number of iterations in power method
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: H405, D205, D400
    if diagonal_loading:
        psd_noise = tik_reg(psd_noise, reg=diag_eps, eps=eps)

    if mode == "power":
        phi = torch.solve(psd_speech, psd_noise)[0]
        e_vec = (
            phi[..., reference_vector, None]
            if isinstance(reference_vector, int)
            else torch.matmul(phi, reference_vector[..., None, :, None])
        )
        for _ in range(iterations - 1):
            e_vec = torch.matmul(phi, e_vec)
            # e_vec = e_vec / complex_norm(e_vec, dim=-1, keepdim=True)
        e_vec = e_vec.squeeze(-1)
    elif mode == "evd":
        # e_vec = generalized_eigenvalue_decomposition(psd_speech, psd_noise)[1][...,-1]
        e_vec = psd_noise.new_zeros(psd_noise.shape[:-1])
        for f in range(psd_noise.shape[-3]):
            try:
                e_vec[..., f, :] = generalized_eigenvalue_decomposition(
                    psd_speech[..., f, :, :], psd_noise[..., f, :, :]
                )[1][..., -1]
            except RuntimeError:
                # port from github.com/fgnt/nn-gev/blob/master/fgnt/beamforming.py#L106
                print(
                    "GEV beamformer: LinAlg error for frequency {}".format(f),
                    flush=True,
                )
                C = psd_noise.size(-1)
                e_vec[..., f, :] = (
                    psd_noise.new_ones(e_vec[..., f, :].shape)
                    / FC.trace(psd_noise[..., f, :, :])
                    * C
                )
    else:
        raise ValueError("Unknown mode: %s" % mode)

    beamforming_vector = e_vec / torch.norm(e_vec, dim=-1, keepdim=True)
    beamforming_vector = gev_phase_correction(beamforming_vector)
    return beamforming_vector
//  [python]
def get_gev_vector(
    psd_noise: torch.Tensor,
    psd_speech: torch.Tensor,
    mode="power",
    reference_vector: Union[int, torch.Tensor] = 0,
    iterations: int = 3,
    diagonal_loading: bool = True,
    diag_eps: float = 1e-7,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Return the generalized eigenvalue (GEV) beamformer vector:

        psd_speech @ h = lambda * psd_noise @ h

    Reference:
        Blind acoustic beamforming based on generalized eigenvalue decomposition;
        E. Warsitz and R. Haeb-Umbach, 2007.

    Args:
        psd_noise (torch.complex64):
            noise covariance matrix (..., F, C, C)
        psd_speech (torch.complex64):
            speech covariance matrix (..., F, C, C)
        mode (str): one of ("power", "evd")
            "power": power method
            "evd": eigenvalue decomposition
        reference_vector (torch.Tensor or int): (..., C) or scalar
        iterations (int): number of iterations in power method
        diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
        diag_eps (float):
        eps (float):
    Returns:
        beamform_vector (torch.complex64): (..., F, C)
    """  # noqa: H405, D205, D400
    if diagonal_loading:
        psd_noise = tik_reg(psd_noise, reg=diag_eps, eps=eps)

    if mode == "power":
        phi = torch.linalg.solve(psd_noise, psd_speech)
        e_vec = (
            phi[..., reference_vector, None]
            if isinstance(reference_vector, int)
            else torch.matmul(phi, reference_vector[..., None, :, None])
        )
        for _ in range(iterations - 1):
            e_vec = torch.matmul(phi, e_vec)
            # e_vec = e_vec / complex_norm(e_vec, dim=-1, keepdim=True)
        e_vec = e_vec.squeeze(-1)
    elif mode == "evd":
        # e_vec = generalized_eigenvalue_decomposition(psd_speech, psd_noise)[1][...,-1]
        e_vec = psd_noise.new_zeros(psd_noise.shape[:-1])
        for f in range(psd_noise.shape[-3]):
            try:
                e_vec[..., f, :] = generalized_eigenvalue_decomposition(
                    psd_speech[..., f, :, :], psd_noise[..., f, :, :]
                )[1][..., -1]
            except RuntimeError:
                # port from github.com/fgnt/nn-gev/blob/master/fgnt/beamforming.py#L106
                print(
                    "GEV beamformer: LinAlg error for frequency {}".format(f),
                    flush=True,
                )
                C = psd_noise.size(-1)
                e_vec[..., f, :] = (
                    psd_noise.new_ones(e_vec[..., f, :].shape)
                    / FC.trace(psd_noise[..., f, :, :])
                    * C
                )
    else:
        raise ValueError("Unknown mode: %s" % mode)

    beamforming_vector = e_vec / torch.norm(e_vec, dim=-1, keepdim=True)
    beamforming_vector = gev_phase_correction(beamforming_vector)
    return beamforming_vector
//  [python]
def Project(self, p, rho, theta):
        # | sin(theta)      cos(theta) | | beta  | = |    x - rho * cos(theta)    |
        # | -cos(theta)     sin(theta) | | gamma |   |    y - rho * sin(theta)    |
        # beta is the distance between the central point (rho * cos(theta), rho * sin(theta)) and the projection
        # gamma is the distance between the point of interest and the projection
        A = torch.tensor([[math.sin(theta), math.cos(theta)], [-math.cos(theta), math.sin(theta)]])
        b = torch.tensor([p[0] - rho * math.cos(theta), p[1] - rho * math.sin(theta)])
        #print(f"DistortionParametersOptimizer.Project(): A.dtype = {A.dtype}; b.dtype = {b.dtype}")
        beta_gamma = torch.solve(b, A)[0]
        p_prime = torch.tensor([rho * math.cos(theta) + beta_gamma[0] * math.sin(theta),
                               rho * math.sin(theta) - beta_gamma[0] * math.cos(theta)])
        return p_prime, beta_gamma[1]
//  [python]
def Project(self, p, rho, theta):
        # | sin(theta)      cos(theta) | | beta  | = |    x - rho * cos(theta)    |
        # | -cos(theta)     sin(theta) | | gamma |   |    y - rho * sin(theta)    |
        # beta is the distance between the central point (rho * cos(theta), rho * sin(theta)) and the projection
        # gamma is the distance between the point of interest and the projection
        A = torch.tensor([[math.sin(theta), math.cos(theta)], [-math.cos(theta), math.sin(theta)]])
        b = torch.tensor([p[0] - rho * math.cos(theta), p[1] - rho * math.sin(theta)])
        #print(f"DistortionParametersOptimizer.Project(): A.dtype = {A.dtype}; b.dtype = {b.dtype}")
        beta_gamma = torch.linalg.solve(A, b)
        p_prime = torch.tensor([rho * math.cos(theta) + beta_gamma[0] * math.sin(theta),
                               rho * math.sin(theta) - beta_gamma[0] * math.cos(theta)])
        return p_prime, beta_gamma[1]
//  [python]
def pade_approximant_inverse(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #return th.linalg.inv(q_sqrt).bmm(p_sqrt)
    return th.solve(p_sqrt,q_sqrt)[0]
//  [python]
def pade_approximant_inverse(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #return th.linalg.inv(q_sqrt).bmm(p_sqrt)
    return th.linalg.solve(p_sqrt,q_sqrt)
//  [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()

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

        # 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)[0] # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().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 = _r2c(self.B) # (H N)
        P = _r2c(self.P)
        Q = P.conj()

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

        # 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, Q_D) # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().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 test_solve_batched_non_contiguous(self, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype).to(device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        self.assertFalse(A.is_contiguous())
        self.assertFalse(b.is_contiguous())
        actual = torch.solve(b, A)[0]
        expected = np.linalg.solve(A.cpu().numpy(), b.cpu().numpy())
        self.assertEqual(actual, expected)
//  [python]
def test_solve_batched_non_contiguous(self, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype).to(device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        self.assertFalse(A.is_contiguous())
        self.assertFalse(b.is_contiguous())
        actual = torch.linalg.solve(A, b)
        expected = np.linalg.solve(A.cpu().numpy(), b.cpu().numpy())
        self.assertEqual(actual, expected)
//  [python]
def phi(self, W,X,T,step):
        """
        Computes the update of the WGD rule

        Args:
            W: particles
            X: inputs training batch
            T: labels training batch

        Return:
            phi: the update to feed the optimizer
            driving force: first term in the update rule
            repulsive force: second term in the update rule
        """
        
        if self.num_train: 
            num_t = self.P.num_train
        else: 
            num_t = 1
            
        W = W.detach().requires_grad_(True)
        
        #computing the driving force
        log_prob = self.P.log_prob(W,X,T)
        score_func = autograd.grad(log_prob.sum(), W)[0]
        
        if self.method == 'kde':
            K_W = self.K(W, W.detach())
            grad_K = autograd.grad(K_W.sum(), W)[0] 
            grad_density = grad_K/ K_W.sum(1,keepdim = True)
            
        elif self.method == 'ssge':
            grad_density = self.pge.compute_score_gradients(W, W)
            
        elif self.method == 'sge':
            eta = 0.01
            K_W = self.K(W, W.detach())
            grad_K = autograd.grad(K_W.sum(), W)[0] 
            K_ = K_W+eta*torch.eye(K_W.shape[0]).to(self.device)  
            grad_density = torch.solve(grad_K,K_)[0]
        
        phi = ( self.ann_schedule[step]*score_func-grad_density) 


        return phi, self.ann_schedule[step]*score_func, grad_density
//  [python]
def phi(self, W,X,T,step):
        """
        Computes the update of the WGD rule

        Args:
            W: particles
            X: inputs training batch
            T: labels training batch

        Return:
            phi: the update to feed the optimizer
            driving force: first term in the update rule
            repulsive force: second term in the update rule
        """
        
        if self.num_train: 
            num_t = self.P.num_train
        else: 
            num_t = 1
            
        W = W.detach().requires_grad_(True)
        
        #computing the driving force
        log_prob = self.P.log_prob(W,X,T)
        score_func = autograd.grad(log_prob.sum(), W)[0]
        
        if self.method == 'kde':
            K_W = self.K(W, W.detach())
            grad_K = autograd.grad(K_W.sum(), W)[0] 
            grad_density = grad_K/ K_W.sum(1,keepdim = True)
            
        elif self.method == 'ssge':
            grad_density = self.pge.compute_score_gradients(W, W)
            
        elif self.method == 'sge':
            eta = 0.01
            K_W = self.K(W, W.detach())
            grad_K = autograd.grad(K_W.sum(), W)[0] 
            K_ = K_W+eta*torch.eye(K_W.shape[0]).to(self.device)  
            grad_density = torch.linalg.solve(K_,grad_K)
        
        phi = ( self.ann_schedule[step]*score_func-grad_density) 


        return phi, self.ann_schedule[step]*score_func, grad_density
//  [python]
def test_solve(self, device, dtype):
        def run_test(n, batch, rhs):
            A_dims = (n, *batch)
            b_dims = (*batch, n, *rhs)
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)

            # Correctness test
            x = torch.solve(b, A)[0]
            if rhs == ():
                Ax = torch.matmul(A, x.unsqueeze(-1))
                Ax.squeeze_(-1)
            else:
                Ax = torch.matmul(A, x)
            self.assertEqual(b.expand_as(Ax), Ax)

            # Check against NumPy
            expected = np.linalg.solve(A.cpu().numpy(), b.expand_as(x).cpu().numpy())
            self.assertEqual(x, expected)

            # Check out= variant
            if rhs == ():
                out = torch.empty_like(x.unsqueeze(-1))
            else:
                out = torch.empty_like(x)
            ans = torch.solve(b, A, out=out)[0]
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

            # Check empty out
            out = torch.empty(0, dtype=dtype, device=device)
            ans = torch.solve(b, A, out=out)[0]
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

        batches = [(), (0, ), (3, ), (2, 3)]
        ns = [0, 5, 32]
        nrhs = [(), (1, ), (5, )]
        for n, batch, rhs in itertools.product(ns, batches, nrhs):
            run_test(n, batch, rhs)
//  [python]
def test_solve(self, device, dtype):
        def run_test(n, batch, rhs):
            A_dims = (n, *batch)
            b_dims = (*batch, n, *rhs)
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)

            # Correctness test
            x = torch.linalg.solve(A, b)
            if rhs == ():
                Ax = torch.matmul(A, x.unsqueeze(-1))
                Ax.squeeze_(-1)
            else:
                Ax = torch.matmul(A, x)
            self.assertEqual(b.expand_as(Ax), Ax)

            # Check against NumPy
            expected = np.linalg.solve(A.cpu().numpy(), b.expand_as(x).cpu().numpy())
            self.assertEqual(x, expected)

            # Check out= variant
            if rhs == ():
                out = torch.empty_like(x.unsqueeze(-1))
            else:
                out = torch.empty_like(x)
            ans = torch.linalg.solve(A, b, out=out)
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

            # Check empty out
            out = torch.empty(0, dtype=dtype, device=device)
            ans = torch.linalg.solve(A, b, out=out)
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

        batches = [(), (0, ), (3, ), (2, 3)]
        ns = [0, 5, 32]
        nrhs = [(), (1, ), (5, )]
        for n, batch, rhs in itertools.product(ns, batches, nrhs):
            run_test(n, batch, rhs)
//  [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 _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()

        # 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 = _r2c(self.B) # (H N)
        P = _r2c(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 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