def __init__(self, nFeatures, args):
        super().__init__()

        nHidden, neq, nineq = 2*nFeatures-1,0,2*nFeatures-2
        assert(neq==0)

        # self.fc1 = nn.Linear(nFeatures, nHidden)
        self.M = Variable(torch.tril(torch.ones(nHidden, nHidden)).cuda())

        Q = 1e-8*torch.eye(nHidden)
        Q[:nFeatures,:nFeatures] = torch.eye(nFeatures)
        self.L = Variable(torch.potrf(Q))

        self.D = Parameter(0.3*torch.randn(nFeatures-1, nFeatures))
        # self.lam = Parameter(20.*torch.ones(1))
        self.h = Variable(torch.zeros(nineq))

        self.nFeatures = nFeatures
        self.nHidden = nHidden
        self.neq = neq
        self.nineq = nineq
        self.args = args 

def get_cov(self, ix=None):

        if ix is None:
            ix = torch.arange(0, self.N)

        return torch.potrf(self.kernel(self.X[ix])
                           + torch.eye(ix.numel())
                                *transform_forward(self.variance),
                           upper=False) 

def pre_factor_kkt(Q, G, A):
    """ Perform all one-time factorizations and cache relevant matrix products"""
    nineq, nz, neq, _ = get_sizes(G, A)

    # S = [ A Q^{-1} A^T        A Q^{-1} G^T           ]
    #     [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]

    U_Q = torch.potrf(Q)
    # partial cholesky of S matrix
    U_S = torch.zeros(neq + nineq, neq + nineq).type_as(Q)

    G_invQ_GT = torch.mm(G, torch.potrs(G.t(), U_Q))
    R = G_invQ_GT
    if neq > 0:
        invQ_AT = torch.potrs(A.t(), U_Q)
        A_invQ_AT = torch.mm(A, invQ_AT)
        G_invQ_AT = torch.mm(G, invQ_AT)

        # TODO: torch.potrf sometimes says the matrix is not PSD but
        # numpy does? I filed an issue at
        # https://github.com/pytorch/pytorch/issues/199
        try:
            U11 = torch.potrf(A_invQ_AT)
        except:
            U11 = torch.Tensor(np.linalg.cholesky(
                A_invQ_AT.cpu().numpy())).type_as(A_invQ_AT)

        # TODO: torch.trtrs is currently not implemented on the GPU
        # and we are using gesv as a workaround.
        U12 = torch.gesv(G_invQ_AT.t(), U11.t())[0]
        U_S[:neq, :neq] = U11
        U_S[:neq, neq:] = U12
        R -= torch.mm(U12.t(), U12)

    return U_Q, U_S, R 

def factor_kkt(U_S, R, d):
    """ Factor the U22 block that we can only do after we know D. """
    nineq = R.size(0)
    U_S[-nineq:, -nineq:] = torch.potrf(R + torch.diag(1 / d.cpu()).type_as(d)) 

def factor_solve_kkt(Q, D, G, A, rx, rs, rz, ry):
    nineq, nz, neq, _ = get_sizes(G, A)

    if neq > 0:
        H_ = torch.cat([torch.cat([Q, torch.zeros(nz, nineq).type_as(Q)], 1),
                        torch.cat([torch.zeros(nineq, nz).type_as(Q), D], 1)], 0)
        A_ = torch.cat([torch.cat([G, torch.eye(nineq).type_as(Q)], 1),
                        torch.cat([A, torch.zeros(neq, nineq).type_as(Q)], 1)], 0)
        g_ = torch.cat([rx, rs], 0)
        h_ = torch.cat([rz, ry], 0)
    else:
        H_ = torch.cat([torch.cat([Q, torch.zeros(nz, nineq).type_as(Q)], 1),
                        torch.cat([torch.zeros(nineq, nz).type_as(Q), D], 1)], 0)
        A_ = torch.cat([G, torch.eye(nineq).type_as(Q)], 1)
        g_ = torch.cat([rx, rs], 0)
        h_ = rz

    U_H_ = torch.potrf(H_)

    invH_A_ = torch.potrs(A_.t(), U_H_)
    invH_g_ = torch.potrs(g_.view(-1, 1), U_H_).view(-1)

    S_ = torch.mm(A_, invH_A_)
    U_S_ = torch.potrf(S_)
    t_ = torch.mv(A_, invH_g_).view(-1, 1) - h_
    w_ = -torch.potrs(t_, U_S_).view(-1)
    v_ = torch.potrs(-g_.view(-1, 1) - torch.mv(A_.t(), w_), U_H_).view(-1)

    return v_[:nz], v_[nz:], w_[:nineq], w_[nineq:] if neq > 0 else None 

def cholesky(self, A, lower=True, warn=False, correct=True):
        return torch.potrf(A, upper=not lower)

    # Tensorflow interface 

def U_UBi_Shb(self, Vs, vs):

        # compute U and V
        V = torch.cat([torch.sqrt(vs[i]) * V for i, V in enumerate(Vs)], 1)
        U = V / torch.sqrt(vs[-1])
        eye = torch.eye(U.shape[1]).cuda()
        B = torch.mm(torch.transpose(U, 0, 1), U) + eye
        # cholB = torch.potrf(B, upper=False)
        # Bi = torch.potri(cholB, upper=False)
        Ub, Shb, Vb = torch.svd(B)
        # Bi = (Vb / Shb).mm(torch.transpose(Vb, 0, 1))
        Bi = torch.inverse(B)
        UBi = torch.mm(U, Bi)

        return U, UBi, Shb 

def __init__(self, nFeatures, args):
        super(OptNet, self).__init__()

        nHidden, neq, nineq = 2*nFeatures-1,0,2*nFeatures-2
        assert(neq==0)

        self.fc1 = nn.Linear(nFeatures, nHidden)
        self.M = Variable(torch.tril(torch.ones(nHidden, nHidden)).cuda())

        if args.tvInit:
            Q = 1e-8*torch.eye(nHidden)
            Q[:nFeatures,:nFeatures] = torch.eye(nFeatures)
            self.L = Parameter(torch.potrf(Q))

            D = torch.zeros(nFeatures-1, nFeatures)
            D[:nFeatures-1,:nFeatures-1] = torch.eye(nFeatures-1)
            D[:nFeatures-1,1:nFeatures] -= torch.eye(nFeatures-1)
            G_ = block((( D, -torch.eye(nFeatures-1)),
                        (-D, -torch.eye(nFeatures-1))))
            self.G = Parameter(G_)
            self.s0 = Parameter(torch.ones(2*nFeatures-2)+1e-6*torch.randn(2*nFeatures-2))
            G_pinv = (G_.t().mm(G_)+1e-5*torch.eye(nHidden)).inverse().mm(G_.t())
            self.z0 = Parameter(-G_pinv.mv(self.s0.data)+1e-6*torch.randn(nHidden))

            lam = 21.21
            W_fc1, b_fc1 = self.fc1.weight, self.fc1.bias
            W_fc1.data[:,:] = 1e-3*torch.randn((2*nFeatures-1, nFeatures))
            # W_fc1.data[:,:] = 0.0
            W_fc1.data[:nFeatures,:nFeatures] += -torch.eye(nFeatures)
            # b_fc1.data[:] = torch.zeros(2*nFeatures-1)
            b_fc1.data[:] = 0.0
            b_fc1.data[nFeatures:2*nFeatures-1] = lam
        else:
            self.L = Parameter(torch.tril(torch.rand(nHidden, nHidden)))
            self.G = Parameter(torch.Tensor(nineq,nHidden).uniform_(-1,1))
            self.z0 = Parameter(torch.zeros(nHidden))
            self.s0 = Parameter(torch.ones(nineq))

        self.nFeatures = nFeatures
        self.nHidden = nHidden
        self.neq = neq
        self.nineq = nineq
        self.args = args