jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve_system(self, rhs, factor, u0, t):
        r"""
        Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.

        Parameters
        ----------
        rhs : dtype_f
            Right-hand side for the linear system.
        factor : float
            Abbrev. for the local stepsize (or any other factor required).
        u0 : dtype_u
            Initial guess for the iterative solver.
        t : float
            Current time (e.g. for time-dependent BCs).

        Returns
        -------
        sol : dtype_u
            The solution of the linear solver.
        """
        solver_type, Id, A, nvars, sol = (
            self.solver_type,
            self.Id,
            self.A,
            self.nvars,
            self.u_init,
        )

        if solver_type == 'direct':
            sol[:] = torch.solve(rhs.flatten(), Id - factor * A)[0].reshape(nvars)
        # TODO: implement torch equivalent of cg
        # elif solver_type == 'CG':
        #     sol[:] = cg(
        #         Id - factor * A,
        #         rhs.flatten(),
        #         x0=u0.flatten(),
        #         tol=lintol,
        #         maxiter=liniter,
        #         atol=0,
        #         callback=self.work_counters[solver_type],
        #     )[0].reshape(nvars)
        else:
            raise ValueError(f'solver type "{solver_type}" not known!')

        return sol
//  [python]
def solve_system(self, rhs, factor, u0, t):
        r"""
        Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.

        Parameters
        ----------
        rhs : dtype_f
            Right-hand side for the linear system.
        factor : float
            Abbrev. for the local stepsize (or any other factor required).
        u0 : dtype_u
            Initial guess for the iterative solver.
        t : float
            Current time (e.g. for time-dependent BCs).

        Returns
        -------
        sol : dtype_u
            The solution of the linear solver.
        """
        solver_type, Id, A, nvars, sol = (
            self.solver_type,
            self.Id,
            self.A,
            self.nvars,
            self.u_init,
        )

        if solver_type == 'direct':
            sol[:] = torch.linalg.solve(Id - factor * A, rhs.flatten()).reshape(nvars)
        # TODO: implement torch equivalent of cg
        # elif solver_type == 'CG':
        #     sol[:] = cg(
        #         Id - factor * A,
        #         rhs.flatten(),
        #         x0=u0.flatten(),
        #         tol=lintol,
        #         maxiter=liniter,
        #         atol=0,
        #         callback=self.work_counters[solver_type],
        #     )[0].reshape(nvars)
        else:
            raise ValueError(f'solver type "{solver_type}" not known!')

        return sol
//  [python]
def rmsd_pbc(batched_pos1, batched_pos2, cell, lig_mask, num_moveable):
    delta_pos = (batched_pos2 - batched_pos1).float() # B x T x 3
    cell = cell.float() # B x 3 x 3
    delta_pos_solve = torch.solve(delta_pos.transpose(-1, -2), cell.transpose(-1, -2))[0].transpose(-1, -2)
    delta_pos_solve[:, :, 0] %= 1.0
    delta_pos_solve[:, :, 0] %= 1.0
    delta_pos_solve[:, :, 1] %= 1.0
    delta_pos_solve[:, :, 1] %= 1.0
    delta_pos_solve[delta_pos_solve > 0.5] -= 1.0
    min_delta_pos = torch.matmul(delta_pos_solve, cell)
    rmsds = torch.sqrt(torch.sum(torch.sum(min_delta_pos ** 2, dim=-1) * lig_mask, dim=-1) / num_moveable)
    return rmsds
//  [python]
def rmsd_pbc(batched_pos1, batched_pos2, cell, lig_mask, num_moveable):
    delta_pos = (batched_pos2 - batched_pos1).float() # B x T x 3
    cell = cell.float() # B x 3 x 3
    delta_pos_solve = torch.linalg.solve(cell.transpose(-1, -2), delta_pos.transpose(-1, -2)).transpose(-1, -2)
    delta_pos_solve[:, :, 0] %= 1.0
    delta_pos_solve[:, :, 0] %= 1.0
    delta_pos_solve[:, :, 1] %= 1.0
    delta_pos_solve[:, :, 1] %= 1.0
    delta_pos_solve[delta_pos_solve > 0.5] -= 1.0
    min_delta_pos = torch.matmul(delta_pos_solve, cell)
    rmsds = torch.sqrt(torch.sum(torch.sum(min_delta_pos ** 2, dim=-1) * lig_mask, dim=-1) / num_moveable)
    return rmsds
//  [python]
def solve_loaded(A: pt.Tensor, b: pt.Tensor, load=1e-6):
    eye = pt.eye(A.shape[-1]).type_as(A)
    load_factor = (
        load * A.abs().max(dim=-2, keepdim=True).values.max(dim=-1, keepdim=True).values
    )
    load_factor = pt.clamp(load_factor, min=load)
    return pt.solve(A + load_factor * eye, b)[0]
//  [python]
def solve_loaded(A: pt.Tensor, b: pt.Tensor, load=1e-6):
    eye = pt.eye(A.shape[-1]).type_as(A)
    load_factor = (
        load * A.abs().max(dim=-2, keepdim=True).values.max(dim=-1, keepdim=True).values
    )
    load_factor = pt.clamp(load_factor, min=load)
    return pt.linalg.solve(A + load_factor * eye, b)
//  [python]
def get_linear_bin_idx(
    cell: torch.Tensor,
    pos: torch.Tensor,
    nbins_s: torch.Tensor,
) -> torch.Tensor:
    """Find the linear bin index of each input pos given a box defined by its cell vectors and a number of bins, contained in the box, for each directions of the box.

    Parameters
    ----------
    cell : torch.Tensor [3, 3]
        cell vectors
    pos : torch.Tensor [-1, 3]
        set of positions
    nbins_s : torch.Tensor [3]
        number of bins in each directions

    Returns
    -------
    torch.Tensor
        linear bin index
    """
    scaled_pos = torch.mm(torch.inverse(cell.t()), pos.t()).t()
    scaled_pos = torch.stack(
        [scaled_pos[:, 0], scaled_pos[:, 1], scaled_pos[:, 2]],
        dim=1,
    )  # redundant op to make torchscript not fail
    bin_index_s = torch.floor(scaled_pos * nbins_s).to(torch.long)
    bin_index_l = ravel_3d(bin_index_s, nbins_s)
    return bin_index_l
//  [python]
def get_linear_bin_idx(
    cell: torch.Tensor,
    pos: torch.Tensor,
    nbins_s: torch.Tensor,
) -> torch.Tensor:
    """Find the linear bin index of each input pos given a box defined by its cell vectors and a number of bins, contained in the box, for each directions of the box.

    Parameters
    ----------
    cell : torch.Tensor [3, 3]
        cell vectors
    pos : torch.Tensor [-1, 3]
        set of positions
    nbins_s : torch.Tensor [3]
        number of bins in each directions

    Returns
    -------
    torch.Tensor
        linear bin index
    """
    scaled_pos = torch.linalg.solve(cell.t(), pos.t()).t()
    scaled_pos = torch.stack(
        [scaled_pos[:, 0], scaled_pos[:, 1], scaled_pos[:, 2]],
        dim=1,
    )  # redundant op to make torchscript not fail
    bin_index_s = torch.floor(scaled_pos * nbins_s).to(torch.long)
    bin_index_l = ravel_3d(bin_index_s, nbins_s)
    return bin_index_l
//  [python]
def dipole_x_beta(self, ptp, x, log_y, z):
        # B x C x 3
        duv = torch.stack(
            [
                x[:, None] - self.local_geom[None, :, 0],
                torch.broadcast_to(
                    torch.exp(2 * log_y)[:, None], ptp.shape
                ),
                z[:, None] - self.local_geom[None, :, 1],
            ],
            axis=2,
        )
        # B x C x 3
        X = duv / (torch.square(duv).sum(axis=2, keepdim=True) + 0.01)
        beta = torch.solve(
            torch.einsum("bck,bc->bk", X, ptp),
            torch.einsum("bci,bcj->bij", X, X),
        )[1]
        return X, beta
//  [python]
def dipole_x_beta(self, ptp, x, log_y, z):
        # B x C x 3
        duv = torch.stack(
            [
                x[:, None] - self.local_geom[None, :, 0],
                torch.broadcast_to(
                    torch.exp(2 * log_y)[:, None], ptp.shape
                ),
                z[:, None] - self.local_geom[None, :, 1],
            ],
            axis=2,
        )
        # B x C x 3
        X = duv / (torch.square(duv).sum(axis=2, keepdim=True) + 0.01)
        beta = torch.linalg.solve(
            torch.einsum("bci,bcj->bij", X, X),
            torch.einsum("bck,bc->bk", X, ptp),
        )
        return X, beta
//  [python]
def state_and_cov_update(Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R):
        H_t = H.transpose(1, 2).double()
        S = (H.bmm(P).bmm(H_t) + R).double()

        Kt = torch.solve(P.bmm(H_t).transpose(1, 2), S)[0].double()
        K = Kt.transpose(1, 2).double()
        dx = bmv(K, r).double()

        dR = torch.zeros(v.shape[0], 3, 3).double()
        dxi = torch.zeros(v.shape[0], 3, 2).double()
        for i in range(0, v.shape[0]):
            dR[i], dxi[i] = IEKF.sen3exp(dx[i, :9])
        dv = dxi[:, :, 0].double()
        dp = dxi[:, :, 1].double()
        Rot_up = dR.bmm(Rot).double()
        v_up = (bmv(dR, v) + dv).double()
        p_up = (bmv(dR, p) + dp).double()

        b_omega_up = b_omega + dx[:, 9:12]
        b_acc_up = b_acc + dx[:, 12:15]

        dR = SO3.exp(dx[:, 15:18]).double()
        Rot_c_i_up = dR.bmm(Rot_c_i).double()
        t_c_i_up = t_c_i + dx[:, 18:21]

        I_KH = IEKF.IdP.double() - K.bmm(H).double()
        P_upprev = I_KH.bmm(P).bmm(I_KH.transpose(1, 2)) + K.bmm(R).bmm(Kt).double()
        P_up = (P_upprev + P_upprev.transpose(1, 2)).double() / 2
        return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
//  [python]
def state_and_cov_update(Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R):
        H_t = H.transpose(1, 2).double()
        S = (H.bmm(P).bmm(H_t) + R).double()

        Kt = torch.linalg.solve(S, P.bmm(H_t).transpose(1, 2)).double()
        K = Kt.transpose(1, 2).double()
        dx = bmv(K, r).double()

        dR = torch.zeros(v.shape[0], 3, 3).double()
        dxi = torch.zeros(v.shape[0], 3, 2).double()
        for i in range(0, v.shape[0]):
            dR[i], dxi[i] = IEKF.sen3exp(dx[i, :9])
        dv = dxi[:, :, 0].double()
        dp = dxi[:, :, 1].double()
        Rot_up = dR.bmm(Rot).double()
        v_up = (bmv(dR, v) + dv).double()
        p_up = (bmv(dR, p) + dp).double()

        b_omega_up = b_omega + dx[:, 9:12]
        b_acc_up = b_acc + dx[:, 12:15]

        dR = SO3.exp(dx[:, 15:18]).double()
        Rot_c_i_up = dR.bmm(Rot_c_i).double()
        t_c_i_up = t_c_i + dx[:, 18:21]

        I_KH = IEKF.IdP.double() - K.bmm(H).double()
        P_upprev = I_KH.bmm(P).bmm(I_KH.transpose(1, 2)) + K.bmm(R).bmm(Kt).double()
        P_up = (P_upprev + P_upprev.transpose(1, 2)).double() / 2
        return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
//  [python]
def solve(self,w,sp):
        dim=len(w)
        a=torch.zeros((self.M,dim),dtype=torch.cfloat)
        for i in range(dim):
            # assemble linear system
            Y=self.K @ sp
            X=1j*(torch.eye(self.M)*w[i]-self.Om)+self.G
            tmp=torch.linalg.lstsq(X,Y,rcond=-1)
            a[:,i]=tmp[0]
#            a[:,i]=torch.linalg.solve(X,Y)
        return a
//  [python]
def solve(self,w,sp):
        dim=len(w)
        a=torch.zeros((self.M,dim),dtype=torch.cfloat)
        for i in range(dim):
            # assemble linear system
            Y=self.K @ sp
            X=1j*(torch.eye(self.M)*w[i]-self.Om)+self.G
#            tmp=np.linalg.lstsq(X,Y,rcond=-1)
#            a[:,i]=tmp[0]
            a[:,i]=torch.linalg.solve(X,Y)
        return a
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        b_inv, _ = torch.solve(eye, b_mat)
        # b_inv = torch.linalg.solve(b_mat, eye)
        # assert torch.allclose(b_inv1, b_inv)
    except:
        b_inv = eye
    return b_inv
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        # b_inv, _ = torch.solve(eye, b_mat)
        b_inv = torch.linalg.solve(b_mat, eye)
        # assert torch.allclose(b_inv1, b_inv)
    except:
        b_inv = eye
    return b_inv
//  [python]
def solve_loaded_general(A, b, load=1e-5, eps=1e-5):

    with pt.no_grad():
        # normalize the rows of A without changing the solution
        # for numerical stability
        norm = pt.linalg.norm(A.detach(), dim=-1, keepdim=True)
        weights = 1.0 / pt.clamp(norm, min=eps)

    load_eye = load * pt.eye(A.shape[-1]).type_as(A)

    # make eigenvalues positive and scale
    A2 = A * weights
    b = b * weights
    A = pt.einsum("...km,...kn->...mn", A2.conj(), A2)
    b = pt.einsum("...km,...kn->...mn", A2.conj(), b)

    return pt.solve(A + load_eye, b)
//  [python]
def solve_loaded_general(A, b, load=1e-5, eps=1e-5):

    with pt.no_grad():
        # normalize the rows of A without changing the solution
        # for numerical stability
        norm = pt.linalg.norm(A.detach(), dim=-1, keepdim=True)
        weights = 1.0 / pt.clamp(norm, min=eps)

    load_eye = load * pt.eye(A.shape[-1]).type_as(A)

    # make eigenvalues positive and scale
    A2 = A * weights
    b = b * weights
    A = pt.einsum("...km,...kn->...mn", A2.conj(), A2)
    b = pt.einsum("...km,...kn->...mn", A2.conj(), b)

    return pt.linalg.solve(A + load_eye, b)
//  [python]
def NewtonScrubStep(weight, grad, hessian, n, l2lambda=0, hessian_device='cpu'):

    original_device = weight.device
    smoothhessian = hessian.to(hessian_device) + (l2lambda)*torch.eye(hessian.shape[0]).to(hessian_device)
    newton = torch.solve(grad.to(hessian_device).unsqueeze(1), smoothhessian)[0].squeeze(1)
    newton = newton.to(original_device)

    new_weight = weight + (1/(n-1))*newton

    return new_weight
//  [python]
def NewtonScrubStep(weight, grad, hessian, n, l2lambda=0, hessian_device='cpu'):

    # smoothhessian = hessian + (l2lambda)*torch.eye(hessian.shape[0]).to(hessian.device)
    # newton = torch.linalg.solve(smoothhessian, grad)

    original_device = weight.device
    smoothhessian = hessian.to(hessian_device) + (l2lambda)*torch.eye(hessian.shape[0]).to(hessian_device)
    newton = torch.linalg.solve(smoothhessian, grad.to(hessian_device))
    newton = newton.to(original_device)

    # removal, towards positive gradient direction
    new_weight = weight + (1/(n-1))*newton

    return new_weight
//  [python]
def CR_NaiveNewton(weight, grad, hessian, l2lambda=0, hessian_device='cpu'):
    original_device = weight.device

    smoothhessian = hessian.to(hessian_device) + (l2lambda)*torch.eye(hessian.shape[0]).to(hessian_device)
    newton = torch.solve(grad.to(hessian_device).unsqueeze(1), smoothhessian)[0].squeeze(1).to(original_device)

    # removal, towards positive gradient direction
    new_weight = weight + newton

    return new_weight
//  [python]
def CR_NaiveNewton(weight, grad, hessian, l2lambda=0, hessian_device='cpu'):
    original_device = weight.device

    smoothhessian = hessian.to(hessian_device) + (l2lambda)*torch.eye(hessian.shape[0]).to(hessian_device)
    newton = torch.linalg.solve(smoothhessian, grad.to(hessian_device)).to(original_device)

    # removal, towards positive gradient direction
    new_weight = weight + newton

    return new_weight
//  [python]
def _project(reference_sources, estimated_source, flen):
    """Least-squares projection of estimated source on the subspace spanned by
    delayed versions of reference sources, with delays between 0 and flen-1
    """

    nsrc, b, nsampl = reference_sources.shape
    def_kw = {'device': reference_sources.device, 'dtype': torch.float64}
    kw = {'device': reference_sources.device, 'dtype': reference_sources.dtype}

    # computing coefficients of least squares problem via FFT ##
    # zero padding and FFT of input data
    reference_sources = torch.cat((reference_sources,
                                   torch.zeros((nsrc, b, flen - 1), **kw)), dim=-1)

    estimated_source = torch.cat((estimated_source, torch.zeros(b, flen - 1, **kw)), dim=-1)
    n_fft = int(2 ** ceil(log2(nsampl + flen - 1.)))
    rs = _fix_shape(reference_sources, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sf = fft(rs)
    es = _fix_shape(estimated_source, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sef = fft(es)
    # inner products between delayed versions of reference_sources
    G = _calc_G(sf, b, nsrc, flen, **kw)
    # inner products between estimated_source and delayed versions of
    # reference_sources

    D = _calc_D(sf, sef, b, nsrc, flen, **kw)
    # Computing projection
    # Distortion filters
    # TODO case in which only few of the Gs are singular but the rest are determined
    if all(torch.det(G) > 0.1):
        C = torch.solve(G, D.unsqueeze(-1))[0].reshape(b, nsrc, flen).permute(0, 2, 1)
    else:
        solutions = []
        for d, g in zip(D, G):
            c = torch.linalg.lstsq(g, d.unsqueeze(1)).solution.reshape(nsrc, flen).T
            solutions.append(c)
        C = torch.stack(solutions)
    # Filtering
    sproj = torch.zeros(b, nsampl + flen - 1, **kw)
    for j in range(b):
        for i in range(nsrc):
            sproj[j, ...] += torch.nn.functional.conv1d(reference_sources[i, j][:-flen + 1][None, None, ...],
                                                        C[j, :, i].flip([0])[None, None, ...],
                                                        padding=flen - 1)[0, 0]

    return sproj
//  [python]
def _project(reference_sources, estimated_source, flen):
    """Least-squares projection of estimated source on the subspace spanned by
    delayed versions of reference sources, with delays between 0 and flen-1
    """

    nsrc, b, nsampl = reference_sources.shape
    def_kw = {'device': reference_sources.device, 'dtype': torch.float64}
    kw = {'device': reference_sources.device, 'dtype': reference_sources.dtype}

    # computing coefficients of least squares problem via FFT ##
    # zero padding and FFT of input data
    reference_sources = torch.cat((reference_sources,
                                   torch.zeros((nsrc, b, flen - 1), **kw)), dim=-1)

    estimated_source = torch.cat((estimated_source, torch.zeros(b, flen - 1, **kw)), dim=-1)
    n_fft = int(2 ** ceil(log2(nsampl + flen - 1.)))
    rs = _fix_shape(reference_sources, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sf = fft(rs)
    es = _fix_shape(estimated_source, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sef = fft(es)
    # inner products between delayed versions of reference_sources
    G = _calc_G(sf, b, nsrc, flen, **kw)
    # inner products between estimated_source and delayed versions of
    # reference_sources

    D = _calc_D(sf, sef, b, nsrc, flen, **kw)
    # Computing projection
    # Distortion filters
    # TODO case in which only few of the Gs are singular but the rest are determined
    if all(torch.det(G) > 0.1):
        C = torch.linalg.solve(G, D.unsqueeze(-1)).reshape(b, nsrc, flen).permute(0, 2, 1)
    else:
        solutions = []
        for d, g in zip(D, G):
            c = torch.linalg.lstsq(g, d.unsqueeze(1)).solution.reshape(nsrc, flen).T
            solutions.append(c)
        C = torch.stack(solutions)
    # Filtering
    sproj = torch.zeros(b, nsampl + flen - 1, **kw)
    for j in range(b):
        for i in range(nsrc):
            sproj[j, ...] += torch.nn.functional.conv1d(reference_sources[i, j][:-flen + 1][None, None, ...],
                                                        C[j, :, i].flip([0])[None, None, ...],
                                                        padding=flen - 1)[0, 0]

    return sproj
//  [python]
def solve_output(self, Hidden_States, y, lambda_reg=1.0):
        '''
        sovle the output with ridge regression, lambda_reg (default) = 1.0
        '''
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)


        col = t_hs.size(1)
        HTH = torch.mm(t_hs.t(), t_hs)
        HTY = torch.mm(t_hs.t(), y)

        I = (lambda_reg * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I

        orig_rank = torch.linalg.matrix_rank(A.float()).item()

        tag = 'Inverse' if orig_rank == col else 'Pseudo-inverse'

        assert tag == 'Inverse'
        
        if tag == 'Inverse':
            W = torch.solve(HTY, A)[0]
        else:
            W = torch.mm(torch.pinverse(A.cpu()),HTY.cpu()).to(self.device)
            
        # self.logger.info('Solving Method: {} \t L2 regular: {}'.format(tag, 'True' if self.lambda_reg != 0 else 'False'))
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y, lambda_reg=1.0):
        '''
        sovle the output with ridge regression, lambda_reg (default) = 1.0
        '''
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)


        col = t_hs.size(1)
        HTH = torch.mm(t_hs.t(), t_hs)
        HTY = torch.mm(t_hs.t(), y)

        I = (lambda_reg * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I

        orig_rank = torch.linalg.matrix_rank(A.float()).item()

        tag = 'Inverse' if orig_rank == col else 'Pseudo-inverse'

        assert tag == 'Inverse'
        
        if tag == 'Inverse':
            W = torch.linalg.solve(A, HTY)
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),HTY.cpu()).to(self.device)
            
        # self.logger.info('Solving Method: {} \t L2 regular: {}'.format(tag, 'True' if self.lambda_reg != 0 else 'False'))
        return W, tag
//  [python]
def solve_linear_regression_(x, y):
        return torch.solve(y, x)[0]
//  [python]
def solve_linear_regression_(x, y):
        # return torch.solve(y, x)[0]
        return torch.linalg.solve(y, x)[0]
//  [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 solve_psd_loaded(A, b, load=1e-5):
    with torch.no_grad():
        # compute normalization weights
        w = torch.diagonal(A, dim1=-2, dim2=-1).detach()
        w = torch.clamp(w, min=0.0)
        w_inv = 1.0 / torch.clamp(torch.sqrt(w), min=1e-5)

    A = (w_inv[..., :, None] * w_inv[..., None, :]) * A
    b = _apply_weights(A, b, w_inv)

    eye = torch.eye(A.shape[-1], device=A.device, dtype=A.dtype)
    A = A + load * eye

    x = torch.solve(b, A)[0]

    x = _apply_weights(A, b, w_inv)

    return x
//  [python]
def solve_psd_loaded(A, b, load=1e-5):
    with torch.no_grad():
        # compute normalization weights
        w = torch.diagonal(A, dim1=-2, dim2=-1).detach()
        w = torch.clamp(w, min=0.0)
        w_inv = 1.0 / torch.clamp(torch.sqrt(w), min=1e-5)

    A = (w_inv[..., :, None] * w_inv[..., None, :]) * A
    b = _apply_weights(A, b, w_inv)

    eye = torch.eye(A.shape[-1], device=A.device, dtype=A.dtype)
    A = A + load * eye

    x = torch.linalg.solve(A, b)

    x = _apply_weights(A, b, w_inv)

    return x
//  [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)

            # 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)

            # 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 exact(g, params, lambd, testing):
    g_flat = tuple_to_vec.tuple_to_vector(g)
    A = derivatives.flat_hessian(g_flat, list(params))

    c = g_flat.detach().to(torch.double)
    A = A.detach().to(torch.double)

    if c.dim() != 1:
        raise ValueError(f"`c` must be a vector, but c = {c}")

    if A.dim() > 2:
        raise ValueError(f"`A` must be a matrix, but A = {A}")

    if c.size()[0] != A.size()[0]:
        raise ValueError("`c` and `A` mush have the same 1st dimension")

    if testing and (A.t() - A).max() > 0.1:
        raise ValueError("`A` is not symmetric")

    h = torch.solve(torch.vstack([-c, torch.zeros_like(c)]), torch.vstack([A + torch.diag(torch.ones_like(c)).mul_(lambd), torch.eye(A.shape[0])]))[0][:A.shape[0]]

    return tuple_to_vec.rollup_vector(h, list(params))
//  [python]
def exact(g, params, lambd, testing):
    g_flat = tuple_to_vec.tuple_to_vector(g)
    A = derivatives.flat_hessian(g_flat, list(params))

    c = g_flat.detach().to(torch.double)
    A = A.detach().to(torch.double)

    if c.dim() != 1:
        raise ValueError(f"`c` must be a vector, but c = {c}")

    if A.dim() > 2:
        raise ValueError(f"`A` must be a matrix, but A = {A}")

    if c.size()[0] != A.size()[0]:
        raise ValueError("`c` and `A` mush have the same 1st dimension")

    if testing and (A.t() - A).max() > 0.1:
        raise ValueError("`A` is not symmetric")

    h = torch.linalg.solve(A + torch.diag(torch.ones_like(c)).mul_(lambd), -c)

    return tuple_to_vec.rollup_vector(h, list(params))
//  [python]
def test_cplx_batchnorm_layer_affine(random_state):
    randn = random_state.randn

    noise = (randn(1000, 10, 17) + 1j * randn(1000, 10, 17)) * 1e-1
    z = randn(1, 10, 1) * (1 + 1j) + noise
    z = np.stack([z.real, z.imag], axis=-1).reshape(*z.shape[:-1], -1)
    x = torch.from_numpy(z).to(torch.float)

    model = torch.nn.Sequential(
        RealToCplx(),
        CplxBatchNorm1d(10, affine=True),
        CplxToReal(),
    )

    # with torch.autograd.detect_anomaly():
    fit(x, x, model, n_steps=5)

    out = predict(x, model).reshape(*x.shape[:-1], -1, 2)
    with torch.no_grad():
        out = out - model[1].bias.reshape(1, 10, 1, 2)
        res = torch.solve(
            out.transpose(-1, -2),
            model[1].weight.permute(2, 0, 1),
        )[1]

    re, im = res[..., 0, :], res[..., 1, :]
    assert np.isclose(float(re.mean()), 0.0, atol=1e-1)
    assert np.isclose(float(im.mean()), 0.0, atol=1e-1)
    assert np.isclose(float((re * re).mean()), 1.0, atol=1e-1)
    assert np.isclose(float((im * im).mean()), 1.0, atol=1e-1)
    assert np.isclose(float((re * im).mean()), 0.0, atol=1e-1)
//  [python]
def test_cplx_batchnorm_layer_affine(random_state):
    randn = random_state.randn

    noise = (randn(1000, 10, 17) + 1j * randn(1000, 10, 17)) * 1e-1
    z = randn(1, 10, 1) * (1 + 1j) + noise
    z = np.stack([z.real, z.imag], axis=-1).reshape(*z.shape[:-1], -1)
    x = torch.from_numpy(z).to(torch.float)

    model = torch.nn.Sequential(
        RealToCplx(),
        CplxBatchNorm1d(10, affine=True),
        CplxToReal(),
    )

    # with torch.autograd.detect_anomaly():
    fit(x, x, model, n_steps=5)

    out = predict(x, model).reshape(*x.shape[:-1], -1, 2)
    with torch.no_grad():
        out = out - model[1].bias.reshape(1, 10, 1, 2)
        res = torch.linalg.solve(
            model[1].weight.permute(2, 0, 1),
            out.transpose(-1, -2),
        )

    re, im = res[..., 0, :], res[..., 1, :]
    assert np.isclose(float(re.mean()), 0.0, atol=1e-1)
    assert np.isclose(float(im.mean()), 0.0, atol=1e-1)
    assert np.isclose(float((re * re).mean()), 1.0, atol=1e-1)
    assert np.isclose(float((im * im).mean()), 1.0, atol=1e-1)
    assert np.isclose(float((re * im).mean()), 0.0, atol=1e-1)
//  [python]
def test_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve

        def run_test(A_dims, B_dims):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            B, A = self.solve_test_helper(A_batch_dims + (A_matrix_size, A_matrix_size), B_dims, device, dtype)
            actual = torch.solve(B, A)[0]
            expected = solve(A.cpu().numpy(), B.cpu().numpy())
            self.assertEqual(actual, expected)

        # test against numpy.linalg.solve
        run_test((5, 5), (2, 0, 5, 3))  # broadcasting with 0 batch dim
        run_test((2, 0, 5, 5), (5, 3))  # broadcasting with 0 batch dim
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting B
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve

        def run_test(A_dims, B_dims):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            B, A = self.solve_test_helper(A_batch_dims + (A_matrix_size, A_matrix_size), B_dims, device, dtype)
            actual = torch.linalg.solve(A, B)
            expected = solve(A.cpu().numpy(), B.cpu().numpy())
            self.assertEqual(actual, expected)

        # test against numpy.linalg.solve
        run_test((5, 5), (2, 0, 5, 3))  # broadcasting with 0 batch dim
        run_test((2, 0, 5, 5), (5, 3))  # broadcasting with 0 batch dim
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting B
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def fit(self, X: torch.Tensor, Y: torch.Tensor) -> ThinPlateSpline:
        """Learn f that matches Y given X

        Args:
            X (Tensor): Control point at source space (X_c)
                Shape: (n_c, d_s)
            Y (Tensor): Control point in the target space (X_t)
                Shape: (n_c, d_t)

        Returns:
            ThinPlateSpline: self
        """
        X = X.to(self.device)
        Y = Y.to(self.device)
        X = _ensure_2d(X)
        Y = _ensure_2d(Y)
        assert X.shape[0] == Y.shape[0]

        n_c, d_s = X.shape
        self.control_points = X

        phi = self._radial_distance(X)

        # Build the linear system AP = Y
        X_p = torch.hstack([torch.ones((n_c, 1), device=self.device), X])

        A = torch.vstack(
            [
                torch.hstack([phi + self.alpha * torch.eye(n_c, device=self.device), X_p]),
                torch.hstack([X_p.T, torch.zeros((d_s + 1, d_s + 1), device=self.device)]),
            ]
        )

        Y = torch.vstack([Y, torch.zeros((d_s + 1, Y.shape[1]), device=self.device)])

        self.parameters = torch.solve(Y, A)[1]  # pylint: disable=not-callable
        self._fitted = True

        return self
//  [python]
def fit(self, X: torch.Tensor, Y: torch.Tensor) -> ThinPlateSpline:
        """Learn f that matches Y given X

        Args:
            X (Tensor): Control point at source space (X_c)
                Shape: (n_c, d_s)
            Y (Tensor): Control point in the target space (X_t)
                Shape: (n_c, d_t)

        Returns:
            ThinPlateSpline: self
        """
        X = X.to(self.device)
        Y = Y.to(self.device)
        X = _ensure_2d(X)
        Y = _ensure_2d(Y)
        assert X.shape[0] == Y.shape[0]

        n_c, d_s = X.shape
        self.control_points = X

        phi = self._radial_distance(X)

        # Build the linear system AP = Y
        X_p = torch.hstack([torch.ones((n_c, 1), device=self.device), X])

        A = torch.vstack(
            [
                torch.hstack([phi + self.alpha * torch.eye(n_c, device=self.device), X_p]),
                torch.hstack([X_p.T, torch.zeros((d_s + 1, d_s + 1), device=self.device)]),
            ]
        )

        Y = torch.vstack([Y, torch.zeros((d_s + 1, Y.shape[1]), device=self.device)])

        self.parameters = torch.linalg.solve(A, Y)  # pylint: disable=not-callable
        self._fitted = True

        return self
//  [python]
def direct_newton(env: DiffEnv, cmd: torch.Tensor, reg_ctrl: float = 0.) -> (torch.Tensor, torch.Tensor):
    """
    Compute a Newton step directly, used to test auto_lin_quad on linear quadratic problem.
    :param env: control environment
    :param cmd: initial command from where to do the Newton step
    :param reg_ctrl: regularization on the control variables, i.e., the command
    :return
        - cmd_opt_newton - optimal command computed by a Newton step
        - opt_newton - optimal value computed by a Newton step
    """
    horizon, dim_ctr = cmd.shape

    # Compute step from a given command
    cmd_flat = deepcopy(cmd.data)
    cmd_flat = cmd_flat.view(-1)
    cmd_flat.requires_grad = True
    cmd_aux = cmd_flat.view(horizon, dim_ctr)

    traj, costs = env.forward(cmd_aux, approx='linquad')
    total_cost = sum(costs)
    feasible = True

    if feasible:
        # Get gradient, hessian, and make a newton step to get the solution
        grad = torch.autograd.grad(total_cost, cmd_flat, create_graph=True)[0]
        # add regularization in the hessian
        tic = time.time()
        hess = auto_multi_grad(grad, cmd_flat) + reg_ctrl * torch.eye(dim_ctr * horizon)
        cmd_opt_newton = - torch.solve(grad.unsqueeze(-1), hess)[0].view(-1)
        print(f'Time Newton: {time.time() - tic}')

        cmd_opt_newton = cmd_opt_newton.view(horizon, dim_ctr)

        traj, costs = env.forward(cmd_aux, approx='linquad')
        lin_dyn_states, lin_dyn_ctrls, \
        quad_cost_states, lin_cost_states, \
        quad_cost_ctrls, lin_cost_ctrls = compute_lin_quad_approx(traj, costs, reg_ctrl)
        lin_quad_env = LinQuadEnv(lin_dyn_states, lin_dyn_ctrls,
                                  quad_cost_states, lin_cost_states,
                                  quad_cost_ctrls, lin_cost_ctrls)
        traj_opt, costs_opt = lin_quad_env.forward(cmd_opt_newton)
        opt_newton = sum(costs_opt)
    else:
        cmd_opt_newton = None
        opt_newton = None
    return cmd_opt_newton, opt_newton
//  [python]
def direct_newton(env: DiffEnv, cmd: torch.Tensor, reg_ctrl: float = 0.) -> (torch.Tensor, torch.Tensor):
    """
    Compute a Newton step directly, used to test auto_lin_quad on linear quadratic problem.
    :param env: control environment
    :param cmd: initial command from where to do the Newton step
    :param reg_ctrl: regularization on the control variables, i.e., the command
    :return
        - cmd_opt_newton - optimal command computed by a Newton step
        - opt_newton - optimal value computed by a Newton step
    """
    horizon, dim_ctr = cmd.shape

    # Compute step from a given command
    cmd_flat = deepcopy(cmd.data)
    cmd_flat = cmd_flat.view(-1)
    cmd_flat.requires_grad = True
    cmd_aux = cmd_flat.view(horizon, dim_ctr)

    traj, costs = env.forward(cmd_aux, approx='linquad')
    total_cost = sum(costs)
    feasible = True

    if feasible:
        # Get gradient, hessian, and make a newton step to get the solution
        grad = torch.autograd.grad(total_cost, cmd_flat, create_graph=True)[0]
        # add regularization in the hessian
        tic = time.time()
        hess = auto_multi_grad(grad, cmd_flat) + reg_ctrl * torch.eye(dim_ctr * horizon)
        cmd_opt_newton = - torch.linalg.solve(hess, grad.unsqueeze(-1), ).view(-1)
        # hess = hess.detach().numpy()
        # grad = grad.unsqueeze(-1).detach().numpy()
        # cmd_opt_newton = cg(hess, grad)[0]
        # cmd_opt_newton = - torch.from_numpy(cmd_opt_newton)
        print(f'Time Newton: {time.time() - tic}')

        cmd_opt_newton = cmd_opt_newton.view(horizon, dim_ctr)

        traj, costs = env.forward(cmd_aux, approx='linquad')
        lin_dyn_states, lin_dyn_ctrls, \
        quad_cost_states, lin_cost_states, \
        quad_cost_ctrls, lin_cost_ctrls = compute_lin_quad_approx(traj, costs, reg_ctrl)
        lin_quad_env = LinQuadEnv(lin_dyn_states, lin_dyn_ctrls,
                                  quad_cost_states, lin_cost_states,
                                  quad_cost_ctrls, lin_cost_ctrls)
        traj_opt, costs_opt = lin_quad_env.forward(cmd_opt_newton)
        opt_newton = sum(costs_opt)
    else:
        cmd_opt_newton = None
        opt_newton = None
    return cmd_opt_newton, opt_newton
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))

    # Scaling step
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)

    # Pade 13 approximation
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R = torch.solve(Q, P)[1]

    # Unsquaring step
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))

    # Scaling step
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)

    # Pade 13 approximation
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R = torch.linalg.solve(P, Q)

    # Unsquaring step
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, 
                       extra_info: Dict[str, torch.Tensor], 
                       theta=None, 
                       mu_r=0, lam_r=0):
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        batch_size = x.shape[0]

        ll_val = 0.0
        for batch in range(batch_size):
            indexed_extra_info = utils.index_dict(extra_info, batch)
            x_proper, u_proper = self._ensure_length(x[batch], u[batch], indexed_extra_info)

            g_hat, g_tilde = self._run_single_jacobian(x_proper, u_proper, theta, 
                                                       extra_info=indexed_extra_info)

            A_batch = self.calc_A(x[batch], u[batch], indexed_extra_info)
            B_batch = self.calc_B(x[batch], u[batch], indexed_extra_info)
            J = self._form_J(A_batch[0], B_batch[0, :-1], indexed_extra_info, print=False)

            g = g_tilde + J @ g_hat
            H = self.hessian(x[batch], u[batch], theta, indexed_extra_info, J)

            ll_val += 0.5 * g @ torch.solve(g, H)[0] + 0.5 * torch.logdet(-H)
        
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val/batch_size
        else:
            return ll_val/batch_size - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, 
                       extra_info: Dict[str, torch.Tensor], 
                       theta=None, 
                       mu_r=0, lam_r=0):
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        batch_size = x.shape[0]

        ll_val = 0.0
        for batch in range(batch_size):
            indexed_extra_info = utils.index_dict(extra_info, batch)
            x_proper, u_proper = self._ensure_length(x[batch], u[batch], indexed_extra_info)

            g_hat, g_tilde = self._run_single_jacobian(x_proper, u_proper, theta, 
                                                       extra_info=indexed_extra_info)

            A_batch = self.calc_A(x[batch], u[batch], indexed_extra_info)
            B_batch = self.calc_B(x[batch], u[batch], indexed_extra_info)
            J = self._form_J(A_batch[0], B_batch[0, :-1], indexed_extra_info, print=False)

            g = g_tilde + J @ g_hat
            H = self.hessian(x[batch], u[batch], theta, indexed_extra_info, J)

            ll_val += 0.5 * g @ torch.linalg.solve(H, g) + 0.5 * torch.logdet(-H)
        
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val/batch_size
        else:
            return ll_val/batch_size - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def woodbury_identity_(invA, U, V):
    # implements the Woodbury identity,
    #
    #   inv(A + U*V) = inv(A) - inv(A)*U*inv(I + V*inv(A)*U)*V*inv(A)
    #
    # with inplace update of invA.
    #
    # Note that using the Woodbury identity multiple times could accumulate numerical erros. 
    invAU = invA @ U
    VinvAU = V @ invAU
    I = torch.eye(VinvAU.shape[0], dtype=VinvAU.dtype, device=VinvAU.device)
    invA.sub_( invAU @ torch.solve(V@invA, I + VinvAU)[0] )
//  [python]
def woodbury_identity_(invA, U, V):
    # implements the Woodbury identity,
    #
    #   inv(A + U*V) = inv(A) - inv(A)*U*inv(I + V*inv(A)*U)*V*inv(A)
    #
    # with inplace update of invA.
    #
    # Note that using the Woodbury identity multiple times could accumulate numerical erros. 
    invAU = invA @ U
    VinvAU = V @ invAU
    I = torch.eye(VinvAU.shape[0], dtype=VinvAU.dtype, device=VinvAU.device)
    invA.sub_( invAU @ torch.linalg.solve(I + VinvAU, V@invA) )
//  [python]
def get_wout(s, sd,sdd, y0s, t,a1s,a0s,fs):
    # y0 = torch.stack([y0 for _ in range(len(s))]).reshape(len(s), -1)


    a0_batch = torch.cat([var_(t) for var_ in a0s], 1)
    a1_batch = torch.cat([var_(t) for var_ in a1s], 1)
    f_batch = torch.cat([var_(t) for var_ in fs], 1)
    WS = []
    for i in range(f_batch.shape[1]):
        y0 = y0s[i,:].reshape(1,-1)
        a0 = a0_batch[:,i].reshape(-1,1)
        a1 = a1_batch[:,i].reshape(-1,1)
        f = f_batch[:,i].reshape(-1,1)

        # print(a0,a1,f)
        DH = (sdd + a1 * sd + a0 * s)
        D0 = f

        h0m = s[0,:].reshape(-1, 1)
        h0d = sd[0,:].reshape(-1, 1)
        W0 = torch.solve(DH.t() @ DH + h0m @ h0m.t() + h0d @ h0d.t(),
                                DH.t() @ D0 + h0m @ (y0[:,0].reshape(1, -1)) + h0d @ (y0[:,1].reshape(1, -1)))[0]
        # print(W0.shape)
        WS.append(W0)
    nWS = (torch.cat(WS)).reshape(f_batch.shape[1],-1)
    return nWS.t()
//  [python]
def get_wout(s, sd,sdd, y0s, t,a1s,a0s,fs):
    # y0 = torch.stack([y0 for _ in range(len(s))]).reshape(len(s), -1)


    a0_batch = torch.cat([var_(t) for var_ in a0s], 1)
    a1_batch = torch.cat([var_(t) for var_ in a1s], 1)
    f_batch = torch.cat([var_(t) for var_ in fs], 1)
    WS = []
    for i in range(f_batch.shape[1]):
        y0 = y0s[i,:].reshape(1,-1)
        a0 = a0_batch[:,i].reshape(-1,1)
        a1 = a1_batch[:,i].reshape(-1,1)
        f = f_batch[:,i].reshape(-1,1)

        # print(a0,a1,f)
        DH = (sdd + a1 * sd + a0 * s)
        D0 = f

        h0m = s[0,:].reshape(-1, 1)
        h0d = sd[0,:].reshape(-1, 1)
        W0 = torch.linalg.solve(DH.t() @ DH + h0m @ h0m.t() + h0d @ h0d.t(),
                                DH.t() @ D0 + h0m @ (y0[:,0].reshape(1, -1)) + h0d @ (y0[:,1].reshape(1, -1)))
        # print(W0.shape)
        WS.append(W0)
    nWS = (torch.cat(WS)).reshape(f_batch.shape[1],-1)
    return nWS.t()