setu.load_estimator#

setu.load_estimator(path, like)[source]#

Load a trained NLE, MixedNLE, or DiscreteNLE from disk.

Uses equinox deserialization. The like parameter provides the pytree structure for deserialization – it must be an estimator with the exact same structure as the saved model. This means:

  • Same architecture (x_dim, theta_dim, flow_type, hidden_dims, n_layers)

  • Same fitted state: if the saved model was trained (has transforms fitted), like must also have fitted transforms. The simplest way to ensure this is to train a model with the same architecture on dummy data, or to keep a copy of the untrained-but-fitted template.

For convenience, use the pattern: save the trained estimator, and when loading, create a new estimator with the same architecture and train it briefly (even 1 epoch) to match the pytree structure.

Parameters:

  • path (str | Path) – File path to load from.

  • like (TypeVar(T)) – An estimator with matching pytree structure. Must have the same hyperparameters and fitted state as the saved model.

Return type:

TypeVar(T)

Returns:

Estimator with weights loaded from disk.

Example

>>> # Save after training
>>> save_estimator(result.nle, "trained_nle.eqx")
>>>
>>> # Load: create matching template (train briefly on dummy data)
>>> template_nle = create_nle(x_dim=3, theta_dim=3, key=key)
>>> template_result = fit_nle(
...     template_nle, dummy_dataset, max_epochs=1, key=key
... )
>>> loaded = load_estimator("trained_nle.eqx", like=template_result.nle)