Transformations¶
relax ¶
relax(
function: Callable[..., PyTree[Array]],
*,
alpha_bb: bool = False,
tighten_bounds: bool = False,
enclosures: bool = False,
) -> Callable[..., Any]
Returns a function that computes McCormick relaxations of function over a box.
function takes an array of variables, then any constant arguments. The returned
function takes a reference point, the lower and upper bounds of the box, and then
those constant arguments. It returns, in the structure of the output of function,
a Relaxation with the interval bounds over the box and the values of the convex
and concave relaxations at the reference point.
The returned function is compiled with jax.jit and composes with jax.vmap,
jax.grad and further jax.jit.
With alpha_bb=True, the result is intersected with the αBB relaxation of
Adjiman et al. (1998), which bounds the Hessian of function over the box with
interval. The result is often tighter for smooth functions on small boxes. It
needs a function that returns a single array, and every primitive in its Hessian
needs a rule.
With tighten_bounds=True, the interval of every intermediate value is narrowed by
the affine relaxations at the middle of the box, after Najman and Mitsos (2019).
Narrower intermediate intervals also tighten the relaxations that depend on them.
This costs about as much as one forward-mode derivative of the relaxation.
With enclosures=True, the returned function takes a flat array of enclosures of
the roots of lax.custom_root after the box, and returns them after the relaxations.
A root enclosed by a box that contains this one needs no search for an enclosure, and
the verified enclosures come back for the boxes inside this one. Use
root_enclosure_size for the length of the array. Entries that are NaN stand for no
enclosure. Only roots at the top level of function take part, and tighten_bounds
cannot be combined with enclosures.
interval ¶
interval(
function: Callable[..., PyTree[Array]],
*,
tighten_bounds: bool = False,
enclosures: bool = False,
) -> Callable[..., Any]
Returns a function that computes interval bounds of function over a box.
The returned function takes the lower and upper bounds of the box, then the
constant arguments of function, and returns the lower and upper bounds of its
output. The convex and concave sides are not needed, so XLA removes them, unless
tighten_bounds, which is passed to relax, needs them.
With enclosures=True, the returned function takes the enclosures of the roots after
the box, as relax does, and returns the verified ones after the bounds.
convex_side ¶
convex_side(
function: Callable[..., Array],
mode: Mode = "forward",
*,
alpha_bb: bool = False,
tighten_bounds: bool = False,
enclosures: bool = False,
) -> Callable[..., tuple[Any, ...]]
Returns a function that computes the convex relaxation and a subgradient.
The returned function takes the same arguments as the one from relax and returns
the value of the convex relaxation of function at the reference point and its
gradient there. function has to return a single array. Computing one side alone
lets XLA remove the computations that only the other side needs. alpha_bb and
tighten_bounds are passed to relax.
With enclosures=True, the function takes the enclosures of the roots after the box
and returns the verified ones as a third result. See relax.
concave_side ¶
concave_side(
function: Callable[..., Array],
mode: Mode = "forward",
*,
alpha_bb: bool = False,
tighten_bounds: bool = False,
enclosures: bool = False,
) -> Callable[..., tuple[Any, ...]]
Returns a function that computes the concave relaxation and a supergradient.
This mirrors convex_side.
affine_relaxation ¶
affine_relaxation(
function: Callable[..., Array],
mode: Mode = "forward",
*,
alpha_bb: bool = False,
tighten_bounds: bool = False,
) -> Callable[..., AffineRelaxation]
Returns a function that computes affine under- and overestimators of function.
The returned function takes the same arguments as the one from relax and returns
an AffineRelaxation. Its planes touch the convex and concave relaxations at the
reference point and bound function on the whole box, since a tangent plane of a
convex function stays below it. function has to return a single array. Both sides
are differentiated in one pass, in forward mode by default. alpha_bb and
tighten_bounds are passed to relax.
AffineRelaxation ¶
Bases: Module
Affine under- and overestimators of a function over a box.
The planes convex + convex_gradient · (x − reference) and
concave + concave_gradient · (x − reference) bound the function from below and
above at every point x of the box.
Mode
module-attribute
¶
The mode of automatic differentiation used for subgradients.