Linrax: JAX-Compatible Simplex Solver

• Linrax is a JAX-native linear programming solver that implements a two-phase tableau simplex method, enabling integration with JAX transformations like jit, vmap, and scan.
• It transforms standard LP constraints into a canonical form while preserving static shapes, using a novel tableau marking strategy to handle dependent constraints.
• Empirical results and control-oriented applications show that Linrax excels in scenarios with many small LP subproblems, especially in robotics, control, and verification workloads.

Searching arXiv for Linrax and closely related software papers. {"query":"linrax JAX simplex linear program solver arXiv", "max_results": 5} {"query":"lrnnx library for Linear RNNs arXiv (Bania et al., 9 Feb 2026)", "max_results": 3} Linrax is a linear programming solver implemented in JAX and designed so that LP solves can be embedded directly inside larger JAX programs while remaining compatible with jit, vectorization, accelerator execution, and automatic differentiation. It is presented as the first simplex-based linear program solver compatible with the JAX ecosystem, with a particular emphasis on robotics, control, and verification workloads in which many relatively small LPs are generated and solved as subroutines inside a compiled pipeline rather than as isolated standalone optimization problems (Gould et al., 23 Sep 2025).

1. Definition, scope, and motivating use cases

Linrax is a JAX-traceable implementation of a two-phase tableau simplex method. Its central design objective is not merely to provide another LP solver, but to make simplex-based LP solution composable with JAX transformations such as jit, vmap, scan, and automatic differentiation. The paper emphasizes that this is especially useful in settings where LPs arise repeatedly inside reachability analysis, robust safety filtering, MPC-like loops, CLF/CBF constructions, and optimization-based policy synthesis (Gould et al., 23 Sep 2025).

The solver is explicitly positioned against two limitations of existing workflows. First, conventional solvers such as SciPy- or Gurobi-based interfaces are not JAX-native, so they interrupt end-to-end compilation and transformation. Second, existing JAX-compatible alternatives are described as relying mainly on first-order methods, which may converge only to moderate accuracy and may struggle or fail on LPs with degenerate or linearly dependent constraints. Linrax is intended to address that gap by combining a simplex method with JAX compatibility and explicit handling of dependent constraints.

The paper repeatedly narrows the intended operating regime. Linrax is especially suited to relatively small LPs solved many times inside a larger JAX computation. It is not presented as a replacement for highly tuned industrial LP engines for large standalone problems. That positioning is important for interpreting both its algorithmic choices and its empirical results.

2. Optimization model and canonicalization

The user-facing LP interface follows a standard split between equality constraints, inequality constraints, and variable bounds. In the paper’s notation, Linrax exposes the general problem

$$\min_{x \in [l, u]} c^\top x \quad \text{s.t.} \quad A_{\mathrm{eq}} x = b_{\mathrm{eq}}, \; A_{\mathrm{ub}} x \le b_{\mathrm{ub}}. \tag{1}$$

Here $x \in \mathbb{R}^n$, $l,u \in \bar{\mathbb{R}}^n$ with $l \le u$, $$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$, $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$, $$A_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}} \times n}$$, $b_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}}}$, and $c \in \mathbb{R}^n$ (Gould et al., 23 Sep 2025).

Algorithmically, Linrax converts this to the canonical LP form

$$\min_{x \ge 0_n} c^\top x \quad \text{s.t.} \quad Ax = b, \tag{2}$$

with $x \in \mathbb{R}^n$0, $x \in \mathbb{R}^n$1, and the standing assumption

$x \in \mathbb{R}^n$2

The paper states that any LP of the general form can be transformed to canonical form, possibly with increased problem dimension. Two transformations are explicitly identified: inequalities are converted to equalities by introducing nonnegative slack variables, and variables on unbounded domains can be represented as differences of two nonnegative variables.

This canonicalization is not merely textbook preprocessing. In Linrax it determines the fixed array shapes used inside JAX tracing. A plausible implication is that the solver’s external simplicity depends on an internal representation disciplined by JAX’s static-shape requirements.

3. Two-phase tableau simplex and the handling of dependent constraints

Linrax uses a two-phase tableau simplex method. Phase I introduces artificial variables $x \in \mathbb{R}^n$3 and solves

$x \in \mathbb{R}^n$4

Because $x \in \mathbb{R}^n$5, the point

$x \in \mathbb{R}^n$6

is immediately feasible for the auxiliary problem, so the artificial variables form the initial basis (Gould et al., 23 Sep 2025).

The phase-I tableau is initialized as

$x \in \mathbb{R}^n$7

The upper block encodes the equality constraints, the last column stores the current basic-variable values, and the bottom two rows represent the original and auxiliary objectives. The initial basis index set is

$x \in \mathbb{R}^n$8

Linrax uses Bland’s rule as its pivot strategy. Negative reduced costs identify entering variables, and the leaving row is chosen by the ratio test

$x \in \mathbb{R}^n$9

If $l,u \in \bar{\mathbb{R}}^n$0, the LP is unbounded. After selecting entering column $l,u \in \bar{\mathbb{R}}^n$1 and leaving row $l,u \in \bar{\mathbb{R}}^n$2, Linrax performs the usual Gauss–Jordan pivot:

$l,u \in \bar{\mathbb{R}}^n$3

$l,u \in \bar{\mathbb{R}}^n$4

The paper’s main algorithmic novelty appears at the transition from phase I to phase II. In conventional simplex implementations, if the LP is feasible but $l,u \in \bar{\mathbb{R}}^n$5 has linearly dependent rows, some artificial variables may remain basic at the end of phase I. Standard solvers can dynamically remove redundant rows or drop artificial-variable structure before phase II, but such value-dependent shape changes are incompatible with JAX tracing. Linrax instead marks the affected rows by replacing them with their elementwise absolute value:

$l,u \in \bar{\mathbb{R}}^n$6

for rows corresponding to artificial variables still in the basis before constructing the phase-II tableau

$l,u \in \bar{\mathbb{R}}^n$7

The paper’s reasoning is that any artificial basic variable remaining after a feasible phase I must have value $l,u \in \bar{\mathbb{R}}^n$8; after marking, any positive coefficient in that row yields ratio $l,u \in \bar{\mathbb{R}}^n$9, so the variable is forced to leave the basis at the earliest available phase-II pivot without any runtime change in tableau shape (Gould et al., 23 Sep 2025).

This marking step is the defining Linrax modification. It translates a standard simplex housekeeping operation into a static-shape, value-only transformation compatible with JAX.

4. JAX compatibility, differentiation, and public interface

Linrax is designed so that the LP solve remains a pure JAX computation. The paper attributes much of its architecture to JAX’s tracing model: intermediate and output arrays must have static shapes known at trace time, and conventional presolve steps such as removing dependent rows, deleting redundant constraints, or changing tableau dimensions after phase I are therefore problematic. Linrax responds by preserving shape throughout the solve and encoding logical changes in values rather than by creating or destroying rows and columns (Gould et al., 23 Sep 2025).

One concrete manifestation of this constraint is the unbounded option. When unbounded=False, Linrax uses the default nonnegative-variable representation. When unbounded=True, variables are represented internally on an unbounded domain, which doubles the number of variables and therefore changes array shapes. For that reason, unbounded is declared as a static argument to jit.

The public interface is deliberately close to SciPy’s linprog style:

$$A_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}} \times n}$$3

The solver returns a SimplexStep, containing x, fun, tableau, and basis, together with a SimplexSolutionType carrying the booleans feasible, bounded, and success. The paper also states that Linrax does not require $l \le u$0 or $l \le u$1 to be full rank.

Automatic differentiation is part of the intended usage, but the paper draws a clear boundary around current capabilities. Linrax is described as compatible with JAX automatic differentiation, and the control case study differentiates a safety metric through repeated LP solves. At the same time, the paper explicitly states that Linrax currently only supports forward-mode autodifferentiation. It does not present custom VJP/JVP rules, implicit differentiation derivations for the KKT system, or a general theory of differentiability across basis changes. This suggests a pragmatic rather than fully formalized AD story: Linrax is differentiable enough to participate in JAX-native pipelines, but not yet specialized for all sensitivity-analysis use cases.

Linrax is open source and available on both PyPI and GitHub, with the package name linrax and repository https://github.com/gtfactslab/linrax.

5. Empirical results and control-oriented applications

The paper reports both isolated solver benchmarks and end-to-end systems benchmarks. The smaller benchmark uses a random LP with $l \le u$2 variables and $l \le u$3 inequality constraints, solved $l \le u$4 times, and reports mean and standard deviation. The larger benchmark places LP solution inside a reachable-set refinement pipeline for the Van der Pol oscillator

$l \le u$5

This second setting is the one the paper treats as representative of Linrax’s intended use case (Gould et al., 23 Sep 2025).

In the random small-LP benchmark, the reported runtimes are:

Solver	JIT time	Runtime
SciPy	—	$l \le u$6 s
Gurobi	—	$l \le u$7 s
CVXPY	—	$l \le u$8 s
JAXopt	$l \le u$9 s	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$0 s
Linrax	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$1 s	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$2 s

These numbers support a narrow conclusion: on a tiny isolated LP, Linrax is slower than the best mature non-JAX standalone solvers, but substantially faster than the compared JAX first-order baseline.

The more consequential results come from LP-based interval refinement for the Van der Pol system. For $$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$3, $$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$4, and $$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$5, the reported end-to-end results are:

Solver	Runtime	Bound size
SciPy	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$6 s	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$7
Gurobi	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$8 s	$$A_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}} \times n}$$9
CVXPY	fail	—
JAXopt	JIT $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$0 s; runtime $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$1 s	$b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$2
Linrax	JIT $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$3 s; runtime $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$4 s	$b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$5

The paper interprets this as the strongest evidence for Linrax’s design philosophy: the advantage arises not because tableau simplex is universally faster in isolation, but because compiling the entire JAX pipeline can dominate the overall workload when many small LPs with dependent constraints are solved inside a larger computation.

A second application concerns safe control filtering for a second-order kinematic bicycle model:

$b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$6

with

$b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$7

and $b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$8. The obstacle is defined by

$b_{\mathrm{eq}} \in \mathbb{R}^{m_{\mathrm{eq}}}$9

and the system is linearized about

$$A_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}} \times n}$$0

The LP subproblems used for reachability refinement are

$$A_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}} \times n}$$1

$$A_{\mathrm{ub}} \in \mathbb{R}^{m_{\mathrm{ub}} \times n}$$2

The paper reports 8 gradient descent steps, 195 s runtime after 52 s of JIT compilation, and remarks that this runtime is likely inflated because only forward-mode autodiff is currently supported. The significance of the example lies less in raw timing than in the fact that Linrax functions as a differentiable LP subroutine inside a nonlinear control-verification loop.

6. Limitations, positioning, and name disambiguation

Linrax’s stated strengths imply equally clear limitations. It is best suited to small LPs, many repeated solves, and JAX-native pipelines. The implementation is a tableau simplex rather than a heavily optimized revised-simplex or barrier method; the paper does not discuss advanced presolve, sparse linear algebra, or warm starts; and its current autodiff support is effectively forward-mode only (Gould et al., 23 Sep 2025).

The paper also avoids claiming a full theory of differentiability for simplex. That omission matters because LP solution maps can be practically delicate near degeneracies, multiple optima, or basis changes. The reported applications show that differentiation through Linrax is useful in practice, but the theoretical and numerical subtleties of differentiating simplex are not treated as solved questions in the source.

A recurrent source of confusion is nomenclature. Linrax is distinct from lrnnx, which is a PyTorch library for linear recurrent neural networks, state-space sequence models, and related LRNN architectures (Bania et al., 9 Feb 2026). It is also distinct from lrux, which is a JAX-based package for low-rank determinant and Pfaffian updates in quantum Monte Carlo workflows (Chen et al., 5 Feb 2026). The similarity in package names can mislead, but the three systems target different problem classes: linear programming in JAX for Linrax, linear recurrent sequence modeling in PyTorch for lrnnx, and structured low-rank matrix updates in JAX for lrux.

Within its intended domain, Linrax is best understood as a JAX-native simplex layer for optimization and control pipelines. Its distinctive technical contribution is the phase-I/phase-II tableau-marking strategy that handles dependent constraints without dynamic shape changes, thereby reconciling simplex with JAX’s compilation model.