sindae.nn_utils
MLP architecture and parameter utilities built on Equinox.
SimpleMLPis the built-in dense feed-forward network used as the unknown term . See Defining a Network Architecture for the requirements it satisfies and how to supply a custom module.flatten_fnandmake_unflatten_fnconvert between the Equinox parameter pytree and a flat 1-D array — the representation used by the decomposition KKT utilities and the simultaneous NLP backend.
Usage¶
import jax
import jax.numpy as jnp
from sindae import SimpleMLP, flatten_fn, make_unflatten_fn
mlp = SimpleMLP(
in_size=2, out_size=1,
widths=[16, 16], # two hidden layers
activations=[jax.nn.softplus, jax.nn.softplus],
key=jax.random.PRNGKey(0),
)
y = mlp(jnp.array([0.5, 1.0])) # forward pass -> shape (1,)
flat = flatten_fn(mlp) # trainable params as a 1-D array
rebuild = make_unflatten_fn(mlp) # inverse: flat array -> SimpleMLP
mlp2 = rebuild(flat) # round-trips back to the same networkAPI reference¶
SimpleMLP¶
class SimpleMLP(
in_size: int,
out_size: int,
widths: List[int],
activations: List[Callable],
*,
key: jax.Array = jax.random.PRNGKey(0),
)Parameters
in_size(int)out_size(int)widths(List[int])activations(List[Callable])key(jax.Array, defaultjax.random.PRNGKey(0))
Methods
__call__¶
__call__(x)Parameters
x
flatten_fn¶
flatten_fn(mlp)Flatten an equinox SimpleMLP’s trainable parameters to a 1-D array.
Parameters
mlp
make_unflatten_fn¶
make_unflatten_fn(mlp)Return a callable that rebuilds a SimpleMLP from a flat parameter array.
Parameters
mlp