tesseract_api.py

# Copyright 2025 Pasteur Labs. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

# Tesseract API module for lp_solve
# Generated by tesseract 0.9.0 on 2025-06-04T13:11:17.208977
import functools
from typing import Self

import jax
import jax.numpy as jnp
import qpax
from pydantic import BaseModel, Field, model_validator

from tesseract_core.runtime import Array, Differentiable, Float32

#
# Schemas
#


class InputSchema(BaseModel):
    Q: Differentiable[Array[(None, None), Float32]] = Field(
        description="Quadratic cost matrix Q for the quadratic program."
    )
    q: Differentiable[Array[(None,), Float32]] | None = Field(
        description="Linear cost vector q for the quadratic program.", default=None
    )
    A: Differentiable[Array[(None, None), Float32]] | None = Field(
        default=None,
        description="Linear equality constraint matrix A for the quadratic program.",
    )
    b: Differentiable[Array[(None,), Float32]] | None = Field(
        default=None,
        description="Linear equality constraint rhs b for the quadratic program.",
    )
    G: Differentiable[Array[(None, None), Float32]] | None = Field(
        description="Linear inequality constraint matrix G for the quadratic program.",
        default=None,
    )
    h: Differentiable[Array[(None,), Float32]] | None = Field(
        description="Linear inequality constraint rhs h for the quadratic program.",
        default=None,
    )
    solver_tol: Float32 | None = Field(
        description="Tolerance for the solver convergence.", default=1e-4
    )
    target_kappa: Float32 | None = Field(
        description="QPAX parameter for QP relaxation.", default=1e-3
    )

    @model_validator(mode="after")
    def validate_shape_inputs(self) -> Self:
        if self.Q.shape[0] != self.Q.shape[1]:
            raise ValueError("Q must be a square matrix.")
        n = self.Q.shape[0]
        # set missing arguments to 'empty' arrays expected by qpax
        if self.q is None:
            self.q = jnp.zeros((n,), dtype=self.Q.dtype)
        if self.A is None:
            self.A = jnp.zeros((0, n), dtype=self.Q.dtype)
        if self.b is None:
            self.b = jnp.zeros((0,), dtype=self.Q.dtype)
        if self.G is None:
            self.G = jnp.zeros((0, n), dtype=self.Q.dtype)
        if self.h is None:
            self.h = jnp.zeros((0,), dtype=self.Q.dtype)
        # check shapes of inputs
        if self.Q.shape[0] != self.q.shape[0]:
            raise ValueError("Q and q must have compatible shapes.")
        if self.A.shape[1] != self.q.shape[0]:
            raise ValueError("A.shape[1] must be equal to number of vars.")
        if self.A.shape[0] != self.b.shape[0]:
            raise ValueError("A and b must have compatible shapes.")
        if self.A.shape[0] >= self.A.shape[1]:
            raise ValueError("A must have fewer rows than columns (m < n).")

        if self.G.shape[1] != self.q.shape[0]:
            raise ValueError("G.shape[1] must be equal to number of vars.")
        if self.G.shape[0] != self.h.shape[0]:
            raise ValueError("G and h must have compatible shapes.")
        return self


class OutputSchema(BaseModel):
    x: Differentiable[Array[(None,), Float32]] = Field(
        description="Optimal solution vector x for the quadratic program."
    )
    objective_value: Float32 = Field(
        description="Optimal objective value for the quadratic program."
    )


jit_solve_qp_primal = jax.jit(
    qpax.solve_qp_primal, static_argnames=("solver_tol", "target_kappa")
)

#
# Required endpoints
#


def apply(inputs: InputSchema) -> OutputSchema:
    x = jit_solve_qp_primal(
        Q=inputs.Q,
        q=inputs.q,
        A=inputs.A,
        b=inputs.b,
        G=inputs.G,
        h=inputs.h,
        solver_tol=inputs.solver_tol,
        target_kappa=inputs.target_kappa,
    )
    obj = jnp.dot(inputs.q, x) + 0.5 * jnp.dot(x, inputs.Q @ x)
    return OutputSchema(
        x=x,
        objective_value=obj,
    )


def vector_jacobian_product(
    inputs: InputSchema,
    vjp_inputs: set[str],
    vjp_outputs: set[str],
    cotangent_vector,
):
    qpsolve_signature = ["Q", "q", "A", "b", "G", "h"]
    # only one output - x
    tangent = cotangent_vector[vjp_outputs.pop()]
    success = False
    kappa = inputs.target_kappa
    max_kappa = 1
    while (not success) and kappa < max_kappa:
        success = True
        solve_fn = functools.partial(
            jit_solve_qp_primal,
            solver_tol=inputs.solver_tol,
            target_kappa=kappa,
        )
        _, f_vjp = jax.vjp(
            solve_fn, inputs.Q, inputs.q, inputs.A, inputs.b, inputs.G, inputs.h
        )
        vjps = f_vjp(tangent)
        out = {}
        for dx in vjp_inputs:
            val = vjps[qpsolve_signature.index(dx)]
            # heuristic attempts to adjust kappa if NaN values are encountered
            # not sure how well this works in practice
            if jnp.any(jnp.isnan(val)):
                kappa *= 2
                success = False
                break
            out[dx] = val
    if not success:
        raise ValueError(
            "Failed to compute VJP with the given inputs. Try increasing target_kappa."
        )
    return out