jaxglobal¶
The top-level package. Importing it turns on 64-bit floats in JAX
(jax_enable_x64) and the on-disk kernel cache.
jaxglobal.minimize ¶
minimize(f, lo, hi, *, ineq=None, eq=None, integer=None, abs_tol=1e-06, rel_tol=1e-06, feas_tol=1e-06, batch=256, batch_sizes=(16,), devices=None, max_iter=5000, n_fw=20, n_kelley=4, local_solve=True, local_solver='auto', n_starts=1, fbbt=True, fbbt_rounds=5, dbbt=True, obbt='root', obbt_rounds=2, branch_point='mid', branch_var='auto', univariate=True, alphabb=True, multivariate=False, lifted=False, quadratic=True, crown=False, cut_pool=128, safe=SAFE_DEFAULT, time_limit=np.inf, verbose=False)
Globally minimize f(x) over lo <= x <= hi, ineq(x) <= 0, eq(x) == 0.
f is any JAX-traceable Python function of a 1-D array returning a
scalar; ineq / eq return a scalar or 1-D array of constraint values.
A point is accepted as feasible when ineq <= feas_tol and |eq| <= feas_tol.
integer marks integer variables (a boolean mask or an array of indices):
their bounds are rounded inward, they are branched on at integer values, and
a returned x has them integral (Result.verified checks it to
INT_TOL). The relaxations are those of the continuous problem.
local_solver picks the upper-bounding local NLP solver: "pounce",
"optax" (pure JAX), "scipy", or "auto" (pounce if installed).
Once there is an incumbent, local searches and integer dives that fail to
improve it are run less often (_LOCAL_BACKOFF, _DIVE_BACKOFF).
The search stops with status "time_limit" after time_limit seconds,
counted from the call: tracing and compiling the bounding kernels count, and
if they are not ready in time the result has no bound (lower_bound=-inf)
and message says so.
Bounds may be infinite: finite ones are derived first from the constraints,
unused variables, and FBBT against a feasible point found by local solves
(domain.py). If some variable stays unbounded the result has status
"unbounded_domain" and message names the variables.
Bound tightening against the incumbent: fbbt (forward-backward interval
propagation over the jaxpr, fbbt_rounds sweeps per node), dbbt
(duality-based range reduction from each node's bound certificate), and
obbt (LP-based tightening of every variable: "root", "tree" for
every node, or False).
Relaxations: univariate replaces the relaxation of every smooth
sub-expression of a single variable by a grid-hull envelope
(univariate.py); alphabb adds whole-function alphaBB with
scaled-Gershgorin alpha from an interval Hessian for smooth functions
(alphabb.py; alphabb="eigen" adds the eigenvalue relaxation of
exact quadratics); multivariate relaxes products by the multivariate
McCormick rule of Tsoukalas & Mitsos, written as up to two extra facets
from its dual (tighter where a factor's interval straddles zero); lifted bounds nodes with the extended-space LP
(lifted.py) instead of the reduced-space cuts. quadratic: a constrained problem whose
functions are all quadratic (and whose tape scalarizes) is bounded, and OBBT-tightened, over
the lifted LP with one McCormick-linked column per product (quadratic.py) instead of the
reduced-space cuts; others are unaffected. crown adds the CROWN
backward linear bound of f (crown.py) to every node bound (the
larger of the two is used): False, True, or "auto" (on when
every primitive of f has a CROWN rule and f has a smooth
saturating activation or ReLU; otherwise off). With True and an
unsupported f the CROWN bound is -inf, i.e. no effect.
A problem whose integer variables are all binary, posed as an epigraph min t, t >= P(x)
with P a polynomial of them at binary points and no other constraint, is solved without a search when the polynomial's induced width is small: exact
variable elimination (pseudoboolean.py) gives the minimum and a minimizer.
safe rounds every quantity feeding a lower bound outward (interval and
McCormick values, envelopes, alphaBB, CROWN lines, LP rows and the dual
certificate) at the primitive-dispatch level; see docs/ROADMAP.md for what
that does and does not guarantee.
After the search the incumbent is re-evaluated outside jit in float64
(box, ineq, eq, f(x) == fun): Result.max_violation and
Result.verified report it, and a failed check changes status
"optimal" to "unverified" (with message) instead of passing silently.
An incumbent repaired to strict feasibility whose gap to the lower bound then exceeds
_REPAIR_GAP times the tolerance is reported with status "unproven".
Branching: branch_point is "mid" (bisect) or "relax" (split at
the relaxation minimizer, pulled toward the midpoint and kept away from the
faces, see _LAM and _BETA); branch_var is "auto" (the default:
"pseudocost" without integer variables, "width" with them; in the 2026-10-20 re-run on
start50 at 60 s pseudocost solves 40, width and strong 38; on int50 width and pseudocost 43 each
and strong 41, see benchmarks/branching_report.py), "width" (largest
relative width among nonlinear variables; on the lifted quadratic LP, relative width
times the McCormick gap of the variable's products at the relaxation point), "pseudocost" (product score
of per-variable bound gains observed so far, on the lifted quadratic LP times a continuous
variable's "width" score; with integer variables, among the integers
fractional at the relaxation point when there are any) or "reliability" (like
pseudocost, but variables with fewer than _ETA observations are
strong-branched, up to _M of them per box, within the batch budget)
or "strong" (every parent is split on its best-scored variable, and the
padding slots of that kernel call strong-branch further candidates, ranked
by the same scores, so it splits as many parents in as many kernel calls
per iteration as "width"; each parent is then split on its _Q best
candidates at once into up to 2**_Q orthants, see _BAL and _ETA_MS).
Result.evaluations counts every box bounded.
Up to batch boxes are bounded per iteration, in kernel calls of the
sizes in batch_sizes (capped by batch; just batch when empty):
a batch is padded to the smallest size that holds it or cut into calls of
the largest. Each size is traced and compiled once, so the default is a
single small size. The kernels compile in background threads and are
cached on disk across processes (jaxglobal/kcache.py; disable with
JAXGLOBAL_CACHE=0). With several JAX devices calls of the largest size
are split across them; on CPU, expose cores by setting
XLA_FLAGS=--xla_force_host_platform_device_count=N before JAX is
imported. devices overrides the device list (e.g.
devices=jax.devices()[:1] to disable splitting).
jaxglobal.Result
dataclass
¶
jaxglobal.make_bounder ¶
make_bounder(f, n_fw=20, univariate=False, alphabb=False, multivariate=False, return_xr=False, safe=False)
Return bound(lo, hi, cutoff) -> (lb, x, f(x), lo', hi') for one box.
[lo', hi'] is the box after duality-based range reduction: every
Frank-Wolfe cut f(z) >= cv(x) + g.(z - x) with f(z) <= cutoff
confines each z_j to an interval. With return_xr a sixth output is
the final Frank-Wolfe iterate (the relaxation minimizer estimate).
jaxglobal.relax ¶
McCormick relaxations of arbitrary JAX functions via jaxpr interpretation.
The user writes an ordinary Python function f(x) with jax.numpy. We
trace it to a jaxpr and re-interpret every primitive on MC objects:
MC(lo, hi, cv, cc)
where, for every point p in the box [lo_box, hi_box] (the "box" of the
original variables):
lo <= value(p) <= hi (interval bounds, constant on the box)
cv(p) <= value(p) <= cc(p) (cv convex in p, cc concave in p)
The input variable is seeded as MC(lo_box, hi_box, p, p). Because every
rule is written with jax.numpy, the relaxation itself is a JAX function:
it can be jit-ed, vmap-ed over many boxes, and differentiated with
jax.grad to get subgradients of the convex relaxation (used for cuts).
Rules come in four kinds:
- structural / positive-linear primitives (slice, reshape, add, reduce_sum, ...) are applied component-wise to (lo, hi, cv, cc);
- signed linear primitives (neg, sub, mul/dot by a constant) swap components by sign;
- bilinear
muluses the McCormick product rule; - univariate primitives use the McCormick composition theorem with either a
hand-written convex/concave envelope (exp, log, sqrt, x**2, |x|, 1/x, ...)
or an automatically derived alphaBB relaxation (sin, cos, tanh, erf, ...)
whose alpha comes from interval bounds on f'' obtained by tracing
jax.grad(jax.grad(primitive))and interpreting that jaxpr in interval mode.
Arithmetic is ordinary floating point. With Ctx.safe every rule's result
is widened outward by a few ulps of its error scale (widen: lo and cv down,
hi and cc up), which makes the bounds valid up to the assumption that XLA's
transcendentals are within a few ulps; see docs/ROADMAP.md for what that
does and does not guarantee.
MC ¶
Bases: NamedTuple
McCormick object: interval [lo, hi] plus convex/concave relaxations.
Unknown ¶
Bases: NamedTuple
A boolean that depends on the variables; we only know where it is decided.
Monomials ¶
Recognizes products of powers c * prod_i b_i ** a_i while an interpreter runs, and
tightens each maximal one with mc_monomial.
observe(name, ins, invals, out) is called after each equation with its operand keys
(jaxpr variables or tape slots; anything hashable, or a literal), the operand values and the
rule's result, and returns the (possibly tightened) result. A record is started by pow with
a constant exponent and extended by mul / div with a positive constant or another
factor (an operand without a record is a factor with exponent 1). Only a product that no later
mul / div consumes (consumed: the keys used as operands of those) and that has at
least two bases, one of them from pow, is tightened: the intermediate products are left to
the product rule, which keeps the kernels small.
LogTerms ¶
Merges like terms a_k * v * log(v) of a sum, while an interpreter runs.
A model read from a file writes a sum such as 9.6 log(x) x + 18.3 log(x) x - 31.0 log(x) x
(ex6_2_6, the Gibbs energy of a mixture): each product is relaxed by itself, so the large
concave terms -31 x log x get the secant, whose gap 31 w / e on [0, w] the search cannot
close near x = 0. As one term (sum_k a_k) x log x the sum is relaxed by mc_xlogx scaled
by the net coefficient (here +1: convex, exact).
observe(name, ins, invals, outvar, out, eqn) is called after each equation (as
Monomials.observe). It records c * log(v) (log, then mul / div by constants,
neg), b * v, products (c log v) * (b v) (an atom A = c b), and sums (add,
sub, neg, mul / div by a constant) of atoms and other values. At a sum that no later
linear operation consumes (consumed: lin_consumed) and that holds an atom, the result is
intersected with rest + sum_v mc_xlogx(v) * A_v. This is the same real function (the
distributive law), so it is a valid relaxation wherever each piece is. With ctx.safe it is
moved outward by 2 SAFE_K ulps of (N + R + 2) times its error scale (N terms with
coefficients rounded R times each, summed in floating point: each rounding is within an ulp
of the scale); mc_xlogx is widened as apply_xlogx widens it and the parts of rest are
the interpreter's (already widened) results.
cvmax ¶
cvmax(a, b)
max(a, b) for a convex relaxation (cv): NaN-propagating like jnp.maximum, but
its derivative is that of the one operand selected by a single comparison.
jnp.maximum's JVP weighs the two operands by two separate comparisons
(x > y / x == y and y > x / y == x). Under jit XLA may compute an operand
more than once (CSE / loop-invariant code motion re-fuse it, and FMA contraction makes the
copies differ in the last bit), so at a near-tie both comparisons can say "the other one":
the weights then sum to 0 (or 2) and jax.grad returns a vector that is not a
subgradient. The Frank-Wolfe and Kelley cuts and duality-based range reduction trust that
gradient, and an invalid cut prunes points with f below the incumbent. One predicate
always picks exactly one operand; at a near-tie both are active up to rounding, so either
gradient gives a valid cut.
safe_mc ¶
safe_mc(x: MC, ok, lo_fill, hi_fill) -> MC
Replace x by a harmless dummy box where ok is False (avoids NaN grads).
concave_rule ¶
concave_rule(f, x: MC, xmax) -> MC
f concave on [x.lo, x.hi] with box maximizer xmax.
alphabb_rule ¶
alphabb_rule(ctx: Ctx, f, x: MC, rng, d2rng) -> MC
Relaxation of a smooth univariate f from bounds on f''.
ucv(z) = f(z) + a/2 (z-lo)(z-hi) is convex and <= f on [lo, hi] when a >= max(0, -min f''); symmetrically for the concave overestimator. These are then composed with x's relaxations via McCormick's theorem.
d2_interval ¶
d2_interval(prim, params, x: MC, safe=False)
Interval bounds of f'' over x's box, for f = the elementwise primitive.
f'' is obtained by autodiff, traced to a jaxpr, and evaluated with this same interpreter in interval-only mode.
mc_mul ¶
mc_mul(a: MC, b: MC, multivariate: bool = False) -> MC
McCormick product rule (Mitsos, Chachuat & Barton 2009 [mitsos2009mccormick]).
mc_log ¶
mc_log(x: MC) -> MC
log (concave increasing). On a box that reaches 0 or below (lo <= 0 < hi) only the
part x > 0 is in log's domain, so no point of the problem has x <= 0 there: the interval
is [-inf, log(hi)] and cc is the concave function that is log on z >= d and its
tangent at d = LOG_DELTA * hi below (finite with a finite slope, above log everywhere);
cv is -inf. Before, the whole result was unbounded, so -log(x - y) on such a box had cv = -inf.
mc_xlogx ¶
mc_xlogx(x: MC) -> MC
Relaxation of x * log(x) (convex on x >= 0, 0 at 0, minimum -1/e at 1/e).
The product rule sees x times log(x), whose interval is unbounded below when the box
touches 0, and gives -inf. Here x * log(x) is relaxed as one convex function of x.
It is defined only for x > 0 (log's domain), so the box is cut to [max(lo, 0), hi]:
a lower end slightly below 0 (FBBT's outward padding) does not make the bound -inf.
- interval:
[g(clip(1/e, L0, U)), max(g(L0), g(U))]withL0 = max(lo, 0); - cv: McCormick composition with
g_d, which isgonz >= dand its tangent atdbelow (d = max(L0, XLOGX_DELTA * min(U, 1))).g_dis convex and finite on the whole line, with finite slopes, so cuts at points withxat (or below) 0 stay finite; it underestimatesgonz >= 0(a tangent of a convex function), by at mostdand only on[0, d); - cc: the secant of
gover[L0, U](its concave envelope there).
xlogx_operand ¶
xlogx_operand(eqn, logs)
Index of v if eqn is v * log(v) (either order), else None.
logs maps the output variable (or tape slot) of every log seen so far to its operand.
xlogx_pair ¶
xlogx_pair(name, ins, logs)
xlogx_operand on the operands ins (jaxpr variables or tape slots) of a name row.
apply_xlogx ¶
apply_xlogx(ctx, eqn, x)
mc_xlogx(x) for the equation eqn (x * log(x)), outward-widened when ctx.safe.
mc_monomial ¶
mc_monomial(coef, factors, safe=False, nround=0)
Relaxation of the monomial coef * prod_i b_i ** a_i (coef > 0) in log space.
On a box where every base is positive and bounded (0 < lo_i, hi_i < inf) the monomial
is coef * exp(L) with L = sum_i a_i log b_i, and McCormick composition gives:
L: intervalsum_i a_i [log lo_i, log hi_i](ordered by the sign ofa_i);cv(L)addsa_i sec_i(cv b_i)fora_i > 0(sec_i: log's secant on[lo_i, hi_i], its convex envelope, which is increasing) anda_i log(cc b_i)fora_i < 0(a_i log bis convex and decreasing inb);cc(L)is the mirror image (a_i log(cc b_i)anda_i sec_i(cv b_i));exp(convex, increasing):cv = coef exp(max(cv L, L.lo))andccthe secant of exp on[L.lo, L.hi]atmin(cc L, L.hi).
cv is convex and cc concave in the point (compositions of convex nondecreasing with convex,
sums with nonnegative weights), and cv <= monomial <= cc on the box. Factor-by-factor McCormick
products lose that a posynomial term is convex in log x; this keeps it, and is intersected with
the product's. With safe everything is moved outward by a relative margin of SAFE_K ulps
times the error scale of the log-space sum (see the comment below), plus nround roundings in
coef (a finite, normal float), and then by an absolute tiny. Returns (ok, MC): ok
marks the boxes where it applies; elsewhere the result is a dummy.
intersect_mc ¶
intersect_mc(a: MC, b: MC) -> MC
Both relaxations at once: the larger lower bounds and the smaller upper bounds.
mul_consumed ¶
mul_consumed(pairs)
Keys used as an operand of a mul or div, from (name, operand keys) pairs.
lin_consumed ¶
lin_consumed(pairs)
Keys used as an operand of a linear operation (add, sub, neg, or mul / div
with a literal other operand), from (name, operand keys) pairs: LogTerms tightens only
sums outside this set.
mc_pow_const ¶
mc_pow_const(x: MC, p) -> MC
x**p for a constant (possibly non-integer) exponent, on x >= 0.
mc_frac ¶
mc_frac(a: MC, b: MC, base: MC) -> MC
Tighten base (McCormick of a * (1/b)) with the convex envelope of a/b.
For a in [aL, aU] with aL >= 0 and b in [bL, bU] with bL > 0 (Tawarmalani & Sahinidis 2001 [tawarmalani2001semidefinite]):
a / b >= ((a + sqrt(aL aU)) / (sqrt(aL) + sqrt(aU)))^2 / b,
since (sqrt(a) - sqrt(aL))(sqrt(a) - sqrt(aU)) <= 0. The right side is
convex (quadratic over linear), nondecreasing in a and nonincreasing in b,
so composing it with a.cv and b.cc keeps it convex.
widen ¶
widen(eqn, invals, out)
Outward-rounded out: lo, cv down and hi, cc up by SAFE_K ulps of the rule's error scale.
The margin is a function of the box only (lo/hi of the result and
of the operands), not of the evaluation point, so cv stays a convex
function of the point with unchanged subgradients and only its value
moves. The error scale is the magnitude of the result plus the sum and the
product of the operand magnitudes (the terms a McCormick facet combines;
their cancellation is what the result's own magnitude would miss), and the
square of a unary operand's (alphaBB shifts by a (z - L)(z - U)).
Sums of n terms use n times the sum of the magnitudes of the terms.
apply_rule ¶
apply_rule(ctx, eqn, invals)
RULES[primitive](ctx, eqn, *invals), outward-widened when ctx.safe.
scan_unrolled ¶
scan_unrolled(eqn)
The scan equation eqn as an equivalent loop-free ClosedJaxpr on its operands.
The body runs length times (in reverse when reverse), threading the carry, reading row
t of each scanned operand and stacking the per-iteration outputs, as export unrolls it.
Interpreting the body once (the inner jaxpr as a plain call) would relax one iteration only.
A scan nested in the body stays a scan and is unrolled when it is interpreted.
inline_call ¶
inline_call(eqn, args, index)
(jaxpr, consts, args) to interpret in place of the higher-order equation eqn, or None
if eqn has no inner jaxpr. index(a) is the concrete value of the operand a (a
cond index) or None if it depends on the variables.
A call (CALLS) is its inner jaxpr; a scan is unrolled (scan_unrolled); a cond
whose index is a constant is its taken branch. Every other higher-order primitive raises
NotImplementedError: its inner jaxpr does not run exactly once on the operands, so
interpreting it as a call would bound the wrong function (one iteration of a loop).
relax ¶
relax(f, multivariate=False, safe=False)
Return rf(x, lo, hi) -> MC: relaxations of f on the box [lo, hi] at x.
multivariate=True relaxes products with the multivariate McCormick rule;
safe=True rounds every result outward (widen).
jaxglobal.interval ¶
interval(f, safe=False)
Return fi(lo, hi) -> (lo, hi): natural interval extension of f.
jaxglobal.MC ¶
Bases: NamedTuple
McCormick object: interval [lo, hi] plus convex/concave relaxations.