def forward(inputs_i, Q, G, A, b, h, U_Q, U_S, R, verbose=False):
    """
    b = A z_0
    h = G z_0 + s_0
    U_Q, U_S, R = pre_factor_kkt(Q, G, A, nineq, neq)
    """
    nineq, nz, neq, _ = get_sizes(G, A)

    # find initial values
    d = torch.ones(nineq).type_as(Q)
    nb = -b if b is not None else None
    factor_kkt(U_S, R, d)
    x, s, z, y = solve_kkt(
        U_Q, d, G, A, U_S,
        inputs_i, torch.zeros(nineq).type_as(Q), -h, nb)
    # x1, s1, z1, y1 = factor_solve_kkt(Q, torch.eye(nineq).type_as(Q), G, A, inputs_i,
    # torch.zeros(nineq).type_as(Q), -h, nb)

    if torch.min(s) < 0:
        s -= torch.min(s) - 1
    if torch.min(z) < 0:
        z -= torch.min(z) - 1

    prev_resid = None
    for i in range(20):
        # affine scaling direction
        rx = (torch.mv(A.t(), y) if neq > 0 else 0.) + \
            torch.mv(G.t(), z) + torch.mv(Q, x) + inputs_i
        rs = z
        rz = torch.mv(G, x) + s - h
        ry = torch.mv(A, x) - b if neq > 0 else torch.Tensor([0.])
        mu = torch.dot(s, z) / nineq
        pri_resid = torch.norm(ry) + torch.norm(rz)
        dual_resid = torch.norm(rx)
        resid = pri_resid + dual_resid + nineq * mu
        d = z / s
        if verbose:
            print(("primal_res = {0:.5g}, dual_res = {1:.5g}, " +
                   "gap = {2:.5g}, kappa(d) = {3:.5g}").format(
                pri_resid, dual_resid, mu, min(d) / max(d)))
        # if (pri_resid < 5e-4 and dual_resid < 5e-4 and mu < 4e-4):
        improved = (prev_resid is None) or (resid < prev_resid + 1e-6)
        if not improved or resid < 1e-6:
            return x, y, z
        prev_resid = resid

        factor_kkt(U_S, R, d)
        dx_aff, ds_aff, dz_aff, dy_aff = solve_kkt(U_Q, d, G, A, U_S,
                                                   rx, rs, rz, ry)
        # D = torch.diag((z/s).cpu()).type_as(Q)
        # dx_aff1, ds_aff1, dz_aff1, dy_aff1 = factor_solve_kkt(Q, D, G, A, rx, rs, rz, ry)

        # compute centering directions
        alpha = min(min(get_step(z, dz_aff), get_step(s, ds_aff)), 1.0)
        sig = (torch.dot(s + alpha * ds_aff, z +
                         alpha * dz_aff) / (torch.dot(s, z)))**3
        dx_cor, ds_cor, dz_cor, dy_cor = solve_kkt(
            U_Q, d, G, A, U_S, torch.zeros(nz).type_as(Q),
            (-mu * sig * torch.ones(nineq).type_as(Q) + ds_aff * dz_aff) / s,
            torch.zeros(nineq).type_as(Q), torch.zeros(neq).type_as(Q))
        # dx_cor, ds_cor, dz_cor, dy_cor = factor_solve_kkt(Q, D, G, A,
        #     torch.zeros(nz).type_as(Q),
        #     (-mu*sig*torch.ones(nineq).type_as(Q) + ds_aff*dz_aff)/s,
        #     torch.zeros(nineq).type_as(Q), torch.zeros(neq).type_as(Q))

        dx = dx_aff + dx_cor
        ds = ds_aff + ds_cor
        dz = dz_aff + dz_cor
        dy = dy_aff + dy_cor if neq > 0 else None
        alpha = min(1.0, 0.999 * min(get_step(s, ds), get_step(z, dz)))
        dx_norm = torch.norm(dx)
        dz_norm = torch.norm(dz)
        if np.isnan(dx_norm) or dx_norm > 1e5 or dz_norm > 1e5:
            # Overflow, return early
            return x, y, z

        x += alpha * dx
        s += alpha * ds
        z += alpha * dz
        y = y + alpha * dy if neq > 0 else None

    return x, y, z
def forward ( inputs_i , Q , G , A , b , h , U_Q , U_S , R , verbose = False ) : nineq , nz , neq , _ = get_sizes ( G , A ) d = torch . ones ( nineq ) . type_as ( Q ) nb = - b if b is not None else None factor_kkt ( U_S , R , d ) x , s , z , y = solve_kkt ( U_Q , d , G , A , U_S , inputs_i , torch . zeros ( nineq ) . type_as ( Q ) , - h , nb ) if torch . min ( s ) < 0 : s -= torch . min ( s ) - 1 if torch . min ( z ) < 0 : z -= torch . min ( z ) - 1 prev_resid = None for i in range ( 20 ) : rx = ( torch . mv ( A . t ( ) , y ) if neq > 0 else 0. ) + torch . mv ( G . t ( ) , z ) + torch . mv ( Q , x ) + inputs_i rs = z rz = torch . mv ( G , x ) + s - h ry = torch . mv ( A , x ) - b if neq > 0 else torch . Tensor ( [ 0. ] ) mu = torch . dot ( s , z ) / nineq pri_resid = torch . norm ( ry ) + torch . norm ( rz ) dual_resid = torch . norm ( rx ) resid = pri_resid + dual_resid + nineq * mu d = z / s if verbose : print ( ( "primal_res = {0:.5g}, dual_res = {1:.5g}, " + "gap = {2:.5g}, kappa(d) = {3:.5g}" ) . format ( pri_resid , dual_resid , mu , min ( d ) / max ( d ) ) ) improved = ( prev_resid is None ) or ( resid < prev_resid + 1e-6 ) if not improved or resid < 1e-6 : return x , y , z prev_resid = resid factor_kkt ( U_S , R , d ) dx_aff , ds_aff , dz_aff , dy_aff = solve_kkt ( U_Q , d , G , A , U_S , rx , rs , rz , ry ) alpha = min ( min ( get_step ( z , dz_aff ) , get_step ( s , ds_aff ) ) , 1.0 ) sig = ( torch . dot ( s + alpha * ds_aff , z + alpha * dz_aff ) / ( torch . dot ( s , z ) ) ) ** 3 dx_cor , ds_cor , dz_cor , dy_cor = solve_kkt ( U_Q , d , G , A , U_S , torch . zeros ( nz ) . type_as ( Q ) , ( - mu * sig * torch . ones ( nineq ) . type_as ( Q ) + ds_aff * dz_aff ) / s , torch . zeros ( nineq ) . type_as ( Q ) , torch . zeros ( neq ) . type_as ( Q ) ) dx = dx_aff + dx_cor ds = ds_aff + ds_cor dz = dz_aff + dz_cor dy = dy_aff + dy_cor if neq > 0 else None alpha = min ( 1.0 , 0.999 * min ( get_step ( s , ds ) , get_step ( z , dz ) ) ) dx_norm = torch . norm ( dx ) dz_norm = torch . norm ( dz ) if np . isnan ( dx_norm ) or dx_norm > 1e5 or dz_norm > 1e5 : return x , y , z x += alpha * dx s += alpha * ds z += alpha * dz y = y + alpha * dy if neq > 0 else None return x , y , z
b = A z_0
    h = G z_0 + s_0
    U_Q, U_S, R = pre_factor_kkt(Q, G, A, nineq, neq)