jaxglobal.export¶

Export a JAX-defined problem to SCIP (PySCIPOpt) or to an AMPL .nl file.

to_scip(f, lo, hi, ineq=..., eq=..., integer=...) builds a pyscipopt.Model; to_nl(f, lo, hi, path, ...) writes a text .nl file that Couenne, BARON, SCIP's AMPL reader and jaxglobal.nl.read_nl can read. Both take the problem in the form minimize takes it: minimize f(x) over lo <= x <= hi, ineq(x) <= 0, eq(x) == 0, with integer a boolean mask or indices. inspect(...) returns the export report without building anything.

How. The functions are traced to a jaxpr, and the jaxpr is evaluated symbolically, element by element, into a hash-consed expression DAG (Graph): each array becomes a numpy object array of node ids, constants are folded, and structural equations (slice, reshape, broadcast, concatenate, gather and dynamic_slice with constant indices, ...) only renumber elements, by binding the primitive on integer index arrays as fbbt_interp.scalarize does. The DAG is then lowered for the target (exact SCIP or .nl operators, exact rewrites, reformulations with auxiliary variables) and written out.

Why not scalarize itself: its rows cover the scalar tapes of .nl models, but it returns None on dot_general, reductions and select_n, and the tape it reads (fbbt.make_tape) inlines a scan body once, as if it were a call. The walker here handles those, unrolls scan / fori_loop and while loops whose condition is a constant, and inlines jit and custom_jvp / custom_vjp calls.

Each operation is translated exactly, rewritten into an identity (tanh(x) = 1 - 2 / (exp(2x) + 1)), reformulated with auxiliary variables (jnp.where, comparisons and sign with binaries and big-M constants from interval bounds over the box), or refused with ExportError naming the primitive. The ExportReport (model.data.report for SCIP, the return value of to_nl) lists which operations took which route, the assumptions made (isnan / isfinite of a model value are taken as False / True: a solver only returns finite points), and what is not exact (approximated):

A reformulated comparison cannot be strict. With strict_eps=0 (the default), at a tie u == v either branch is allowed, so the model is a relaxation where the branches of a jnp.where (or the values of sign) differ at the tie: a piecewise-constant function can then reach a value at a breakpoint that JAX never produces. strict_eps > 0 makes the comparison strict by excluding 0 < |u - v| < strict_eps on the wrong side, a restriction instead. Either way, evaluate f with JAX at the returned x.
reduce_precision and float64 -> float32 casts are exported as identities (the solvers compute in double precision).

Refused rather than exported wrongly: a jnp.where branch (or comparison operand) that is undefined on part of the box, such as sqrt(x) with x < 0 possible, because the reformulation evaluates both branches everywhere and would cut that part of the box off (write jnp.sqrt(jnp.maximum(x, 0))); big-M constants of 1e20 or more (SCIP's infinity; above 1e8 they are reported and warned about, ExportBigMWarning); division by the constant 0; an exported intermediate whose interval bound reaches 1e20 over the box while its operands are bounded (SCIP would treat those points as invalid evaluations and cut them off). The tanh, sigmoid, asinh and acosh rewrites choose overflow-free forms from the argument's interval so they do not hit this.

ExportError ¶

Bases: ValueError

The problem uses an operation the target cannot represent exactly.

ExportSizeWarning ¶

Bases: UserWarning

The exported model is larger than max_size scalar operations.

ExportBigMWarning ¶

Bases: UserWarning

A big-M constant above BIGM_WARN: the reformulation is numerically weak.

Graph ¶

Hash-consed DAG; node i is (op, args, param) with args ids below i.

op is "var" (param: index), "const" (param: value), an operator of EVAL, "sum" (n-ary) or "unsupported" (param: message). select(p, f, t) is t where p else f (lax.select_n order).

op ¶

op(name, *args)

Node name(*args), constant-folded and lightly simplified.

ExportReport `dataclass` ¶

What the export did with each operation of the model.

exact: operations written as the target's own operator. rewritten: operations written as an exact identity in other operators (how). reformulated: operations replaced by auxiliary variables and constraints (how). refused: operations the target cannot represent (to_* raise ExportError). approximated: operations whose export is not exact (how): comparisons and sign (relaxed at ties, or restricted near them with strict_eps), reduce_precision and float64 -> float32 casts (identities). When it is not empty, evaluate f with JAX at the solver's x before trusting the objective value.

SCIPExport `dataclass` ¶

Attached to the model as model.data.

inspect ¶

inspect(f, lo, hi, *, ineq=None, eq=None, integer=None, target='scip', minmax='abs', nonsmooth='native', portable=False, strict_eps=0.0)

The ExportReport of exporting to target ("scip" or "nl"), without raising on refusals.

to_scip ¶

to_scip(f, lo, hi, *, ineq=None, eq=None, integer=None, name='jaxglobal', minmax='abs', max_size=100000.0, share_min=50, strict_eps=0.0)

A pyscipopt.Model for: minimize f(x) s.t. ineq(x) <= 0, eq(x) == 0, lo <= x <= hi.

The arguments are those of jaxglobal.minimize. model.data is a SCIPExport with the variables of x and the ExportReport; solution(model) returns the best solution as an array.

minmax="abs" writes max / min (and relu, clip) as the exact identities (a + b +- |a - b|) / 2; "bigm" uses a binary and big-M constraints instead. jnp.where and sign always use binaries with big-M constants from interval bounds over the box, and are refused when those are not finite. Reformulated comparisons (and sign) are relaxed at ties unless strict_eps > 0 (see the module docstring and report.approximated). A shared subexpression with at least share_min operations becomes an auxiliary variable (w == expr), so unrolled loops do not grow into exponentially large expression trees, and a product operand becomes one before pyscipopt would expand the product past MAX_TERMS monomials. A nonlinear objective is moved into a constraint f(x) <= t with minimize t. Warns (ExportSizeWarning) above max_size scalar operations and (ExportBigMWarning) for big-M constants above BIGM_WARN. Raises ExportError naming the operation when something cannot be exported.

solution ¶

solution(model, sol=None)

x of the best solution of a model built by to_scip (or of sol), as an array.

to_nl ¶

to_nl(f, lo, hi, path, *, ineq=None, eq=None, integer=None, nonsmooth='native', minmax='abs', x0=None, max_size=100000.0, portable=False, strict_eps=0.0, share_min=50)

Write the problem (arguments as for minimize) to the text .nl file path.

nonsmooth="native" keeps max, min, jnp.where (if) and comparisons as .nl operators: exact, no extra variables, read by jaxglobal.nl.read_nl and ASL-based solvers that implement them. "reformulate" replaces them with binaries and big-M constraints as to_scip does (strict_eps as there).

Which operators a reader accepts varies. SCIP's .nl reader (scip binary) takes + - * / neg pow exp log sqrt abs sin cos and n-ary sums only; Couenne also takes tan tanh sinh cosh. Neither takes min / max, if and comparisons, atan asin acos atan2 or the inverse hyperbolic functions. portable=True writes only operators both accept: the others are rewritten as for to_scip (tanh with exp, ...), nonsmooth ones are reformulated (nonsmooth is ignored), and inverse trigonometric functions are refused. Variables are written continuous first, then integer (the order ASL requires); the report's var_order[k] is the model variable of .nl variable k (the first n_vars model variables are x, the rest auxiliary). Shared subexpressions become defined variables (V segments); with portable=True there are none (Couenne 0.5.8 misreads V segments, even in files AMPL writes; use portable=True with Couenne), and a shared subexpression with at least share_min operations becomes an auxiliary variable with w == expr instead. Returns the ExportReport.