Machine Learning Manager

MLManager owns the neural network, optimizer, and training state. It is constructed automatically by NNFE.from_config but can also be used standalone for generic JAX/Equinox training workflows.

MLManager

nnfe.ml.MLManager

Manages the neural network, optimizer, and training state for NNFE.

On construction this class:

1. Builds or loads the network(s) from :mod:nnfe.networks.
2. Initialises (or loads) the weights.
3. Builds the Optax optimizer and its learning-rate schedule.
4. Initialises the optimizer state.

Parameters:

Name	Type	Description	Default
MLConfig		An :class:~nnfe.ml_config.MLConfig instance describing the networks, optimizer, and training hyperparameters.	required
**kwargs		Must include out_size (int) when any network specifies out_size: "dofs" in its config.	required

Attributes:

Name	Type	Description
config		The :class:~nnfe.ml_config.MLConfig used to build this manager.
network		The Equinox model (single or combined via :class:~nnfe.networks.Sum_models).
filter		A pytree mirror of network whose leaves are True for trainable parameters and False for frozen ones.
optimizer		The Optax optimizer.
lr_scheduler		The Optax learning-rate schedule callable.
opt_state		Current Optax optimizer state.
epochs		Total training epochs (cast to int).
batch_size		Batch fraction (cast to int after scaling by the sampler size in :class:~nnfe.nnfe_object.NNFE).

from_config(MLConfig, **kwargs) classmethod

from_yaml(param_file, **kwargs) classmethod

create_network(params, out_size, key=0)

Build the full network (or combined network) from config.

Iterates over all network configs, instantiates each via :meth:network_from_config, and combines them with :class:~nnfe.networks.Sum_models if more than one is present. Also builds the trainability filter, freezing any networks marked static=True.

Parameters:

Name	Type	Description	Default
params		Ordered dict of :class:~nnfe.ml_config.NetworkConfig objects (from MLConfig.networks).	required
out_size		Output dimension (number of DoFs).	required
key		JAX PRNG key.	0

Returns:

Type	Description
	Tuple of (network, filter) where filter is a pytree of
	booleans indicating trainable parameters.

create_optimizer(OptimizerConfig)

Build the Optax optimizer and learning-rate schedule.

Constructs a piecewise :func:optax.join_schedules from the scheduler sub-config and injects it as the learning_rate argument of the optimizer.

Parameters:

Name	Type	Description	Default
OptimizerConfig		An :class:~nnfe.ml_config.OptimizerConfig instance.	required

Returns:

Type	Description
	Tuple of (optimizer, scheduler) where scheduler is the
	callable learning-rate schedule.

network_from_config(NetworkConfig, **kwargs)

Instantiate a single network from a :class:~nnfe.ml_config.NetworkConfig.

Resolves the "dofs" placeholder for out_size, injects the activation function from :mod:jax.nn, constructs the network, and either initialises or loads its weights. Prints the total parameter count.

Parameters:

Name	Type	Description	Default
NetworkConfig		A :class:~nnfe.ml_config.NetworkConfig instance.	required
**kwargs		Must include out_size (int) and key (:class:jax.random.PRNGKey).	required

Returns:

Type	Description
	An initialised Equinox module.

init_linear_weight(model, key)

Re-initialise all eqx.nn.Linear weights and biases in model.

Replaces every weight matrix and bias vector in model with values drawn from :meth:trunc_weight / :meth:trunc_bias. If the model has no biases the bias replacement is silently skipped.

Parameters:

Name	Type	Description	Default
model		An Equinox module whose Linear layers should be re-initialised.	required
key		JAX PRNG key used to generate replacement values.	required

Returns:

Type	Description
	A new Equinox module with re-initialised weights (and biases if
	present).

load_network(network, path)

Load serialised Equinox leaves from path into network.

Attempts a direct deserialisation first. If the on-disk weights are stored in a different floating-point dtype (e.g. float32 vs float64), the model is cast to float32 before retrying.

Parameters:

Name	Type	Description	Default
network		An Equinox module whose structure matches the serialised file.	required
path		Path to the .eqx file produced by :func:equinox.tree_serialise_leaves.	required

Returns:

Type	Description
	The same module with its leaves replaced by the loaded weights.

filtering(filter, model_ind)

Mark the sub-model at index model_ind as non-trainable in filter.

Sets the corresponding subtree of the filter pytree to False so that Equinox excludes it from gradient computation.

Parameters:

Name	Type	Description	Default
filter		A pytree (mirroring network) of boolean leaves.	required
model_ind		Index into network.models to freeze.	required

Returns:

Type	Description
	Updated filter pytree.

trunc_weight(weight: Array, key: PRNGKey) -> Array

Return a new weight matrix sampled from a truncated normal.

Values are drawn from TruncNormal(0, 1e-6) clipped to [-1, 1], giving a near-zero initialisation suitable for residual minimisation.

Parameters:

Name	Type	Description	Default
weight	Array	Existing weight array (used only for shape).	required
key	PRNGKey	JAX PRNG key.	required

Returns:

Type	Description
Array	New weight array with the same shape as weight.

trunc_bias(bias: Array, key: PRNGKey) -> Array

Return a new bias vector sampled from a truncated normal.

Parameters:

Name	Type	Description	Default
bias	Array	Existing bias array (used only for shape).	required
key	PRNGKey	JAX PRNG key.	required

Returns:

Type	Description
Array	New bias array with the same shape as bias.

dump_config(save_dir: Path, filename: str)

Dumps this specific manager's configuration to a YAML file.