Rules¶

register_rule ¶

register_rule(
    primitive: Primitive | str, rule: RelaxationRule
) -> RelaxationRule

Registers rule as the relaxation rule of primitive and returns the rule.

A rule that is already registered for the primitive is replaced.

RelaxationRule ¶

Bases: Protocol

A function that relaxes one primitive.

The interpreter calls a rule only when at least one operand is a Relaxation. A rule returns one operand, or a sequence of them for a primitive with multiple results.

RULES `module-attribute` ¶

RULES: dict[str, RelaxationRule] = {}

Relaxation rules keyed by primitive name.

Operand `module-attribute` ¶

Operand: TypeAlias = Relaxation | Array

A primitive operand, relaxed when it depends on the variables and plain otherwise.

Params `module-attribute` ¶

Params: TypeAlias = Mapping[str, Any]

The parameters of a primitive.

Their types depend on the primitive. They hold shapes, dtypes, axes, jaxprs and more, so no type narrower than Any covers every value.

Building blocks¶

as_relaxation ¶

as_relaxation(operand: Operand) -> Relaxation

Returns the operand as a Relaxation, treating plain arrays as exact constants.

compose_convex ¶

compose_convex(
    operand: Relaxation,
    function: UnivariateFunction,
    argmin: Array,
) -> Relaxation

Relaxes function(operand) for a function that is convex on the interval.

argmin is the point of the interval where function is smallest. The convex side uses the function itself and the concave side uses its secant.

compose_concave ¶

compose_concave(
    operand: Relaxation,
    function: UnivariateFunction,
    argmax: Array,
) -> Relaxation

Relaxes function(operand) for a function that is concave on the interval.

This mirrors compose_convex, with the secant on the convex side and the function itself on the concave side.

compose ¶

compose(
    operand: Relaxation,
    convex_envelope: Envelope,
    argmin: Array,
    concave_envelope: Envelope,
    argmax: Array,
    lower: Array,
    upper: Array,
) -> Relaxation

Relaxes f(operand) with McCormick's composition rule.

convex_envelope and concave_envelope bound f from below and above on the interval of the operand. argmin is a point where the convex one is smallest and argmax a point where the concave one is largest. lower and upper bound f on the interval.

mid ¶

mid(
    convex: Array, concave: Array, point: ArrayLike
) -> Array

Returns McCormick's mid operator, the median of convex, concave and point.

The median is written as an explicit choice of branch. A median built from jnp.minimum and jnp.maximum has the same value, but where convex and concave tie, jax.grad splits the gradient between them. That mixes a subgradient of the convex side with a supergradient of the concave side, and the resulting tangent planes can cut off feasible points. Ties are common on the boundary of the box, which is where the solutions of linear programs lie.

Rounding can leave concave an ulp below convex. The branches therefore compare point against the smaller and the larger side, and a point between the two sides is returned as is. The envelope is then constant at its minimum or maximum, which still bounds the function.

secant ¶

secant(
    value_at_lower: Array,
    value_at_upper: Array,
    lower: Array,
    upper: Array,
    point: Array,
) -> Array

Returns the line through (lower, value_at_lower) and (upper, value_at_upper).

The line is evaluated at point. On an interval of zero width it is the constant value_at_lower.