cottax.relational.consumers.execution.jax_solvers

The jax backend of the driver contracts: Newton, Picard, BFGS and SLSQP, each chosen per contract in ImplementDrivers.

Dependency:

canonical_nlp (canonical, box: relational.programme computed over the

||—————– values at the condition’s reads; next_values, Picard’s || step; NotSquare) / /

optimistix scipy_slsqp (optimistix: Newton for a RootFinder, Picard for a

FixedPointIterator, BFGS for a Minimiser; scipy_slsqp: SLSQP for a ConstrainedMinimiser, through pure_callback)

class Newton(rtol: float = 0.0001, atol: float = 0.0001, max_steps: int = 256)[source]

Bases: Optimistix[RootFinder]

optx.Newton on the equality residuals, as many numbers as the unknowns.

solve[V: Place](posed: Posed[V], context: Context, data: ByRole)[source]

The optimistix solve answering posed: its function takes args=(posed, context).

Return type:

Solution

implements

alias of RootFinder

class Picard(rtol: float = 0.0001, atol: float = 0.0001, max_steps: int = 256)[source]

Bases: Optimistix[FixedPointIterator]

optx.FixedPointIteration: u <- g(u), no Jacobian taken.

solve[V: Place](posed: Posed[V], context: Context, data: ByRole)[source]

The optimistix solve answering posed: its function takes args=(posed, context).

Return type:

Solution

implements

alias of FixedPointIterator

class BFGS(rtol: float = 0.0001, atol: float = 0.0001, max_steps: int = 256)[source]

Bases: Optimistix[Minimiser]

optx.BFGS on the one objective.

solve[V: Place](posed: Posed[V], context: Context, data: ByRole)[source]

The optimistix solve answering posed: its function takes args=(posed, context).

Return type:

Solution

implements

alias of Minimiser

class SLSQP(tol: float = 1e-08, max_steps: int = 200, strict: bool = True, jacobian: Callable = <function jacfwd>)[source]

Bases: DriverImplementation[ConstrainedMinimiser]

scipy.optimize.minimize(method=’SLSQP’) on the canonical form: the objective, the equalities as eq, the inequalities negated as ineq (scipy’s >= 0), and the box as scipy’s bounds, so every iterate stays inside it.

needs: AbstractClassVar[frozenset[type[Role]]] = frozenset({<class 'cottax.core.driven.roles.Start'>})
provides: AbstractClassVar[frozenset[type[Role]]] = frozenset({<class 'cottax.core.driven.roles.Steps'>})
tol: float = 1e-08
max_steps: int = 200
strict: bool = True

Raise when scipy does not report success; off, the answer is taken as it stands.

jacobian(argnums: int | Sequence[int] = 0, has_aux: bool = False, holomorphic: bool = False)

How the Jacobian is taken: jax.jacfwd, jax.jacrev, or any hashable function of a function, such as one pulling rows back one at a time to save memory.

Return type:

Callable

implements

alias of ConstrainedMinimiser

Submodules