Potential functions and transformations

dynax.ConvexNormalization

Normalization for convex scalar potentials.

This module normalizes a given function \(f:\mathbb{R}^n \times \mathcal{P} \rightarrow \mathbb{R},\, (x; \text{args}) \mapsto f(x; \text{args})\) convex in \(x\), by applying the following transformation: $$ F(x; \text{args}) = f(x; \text{args}) - f(x^\ast; \text{args}) - \frac{\partial f}{\partial x} \bigg\vert_{x^\ast}\cdot (x-x^\ast) + \epsilon\left\lVert x-x^\ast \right\rVert^2. $$ with \(x^\ast \in \mathbb{R}^n\) and \(\epsilon \geq 0\). This ensures the resulting function \(F(x; \text{args})\) has the following properties, given \(f\) is convex in \(x\):

• \(F\) is convex in \(x\),
• \(F\) is strictly convex in \(x\) if \(\epsilon>0\),
• \(F\) has a unique minimum at \(x^\ast\) and,
• \(F(x^\ast) = 0\).

The intended use of this module is to normalize a (partially) convex neural network and use the resulting function as a Lyapunov candidate function, or as scalar potential to ensure global Lyapunov stability by construction.

__init__(func, state_size, argmin_init=zeros, argmin_learnable=False, epsilon=1e-06, dtype=None, *, key)

Initialize the Lyapunov normalization.

PARAMETER	DESCRIPTION
func	Function f(x, args) -> Scalar, convex in x. TYPE: Callable[[Array, PyTree], Shaped[Array, '']]
state_size	State size n needed to determine the size for the minimizer x*. TYPE: int
argmin_init	JAX initializer for the minimum location x*. Can be any function with signature (key, shape, dtype) -> Array but typically is a JAX initializer. TYPE: Initializer DEFAULT: zeros
argmin_learnable	If True, the minimizer x* is learnable. Otherwise, it is wrapped in klax.NonTrainable, preventing updates during optimization. TYPE: bool DEFAULT: False
epsilon	Small value to ensure the resulting function is strictly convex. Warning If epsilon is an inexact array (eqx.is_inexact_array(epsilon)) == True) then it will receive gradient updates. Make sure to use a plain python float or wrap it in klax.non_trainable. TYPE: float DEFAULT: 1e-06
dtype	The dtype to use for the minimum. Defaults to either jax.numpy.float32 or jax.numpy.float64 depending on whether JAX is in 64-bit mode. TYPE: type | None DEFAULT: None
key	PRNG key for random initialization of the minimum. TYPE: PRNGKeyArray

__call__(x, args)

Evaluate the normalized function.

PARAMETER	DESCRIPTION
x	Evaluation point for the normalized function. This should be a vector of shape (n,). TYPE: Array
args	Additional arguments passed to the wrapped function. TYPE: PyTree

RETURNS	DESCRIPTION
Shaped[Array, '']	Value of the normalized function at x, args.

dynax.NoArgsPotential

Wraps a callable x -> Scalar for use with ConvexNormalization.

Convenience wrapper that just adds the args argument to the call signature to ensure compatibility.

Example

func = klax.nn.FICNN(4, "scalar", ...)
potential = ConvexNormalization(NoArgsPotential(func), state_size=4)

__init__(func)

Initialize self. See help(type(self)) for accurate signature.

__call__(x, args)