record_fit

nyx.infer.convergence.record_fit(opt, model, probes=None, *, max_steps=200, stride=1, jit=True, callback=None, stop_on_diverged=True)[source]

Step opt over model, recording the loss and the probes.

Every stride-th step is recorded, always including step 0 and the final one.

Parameters:

  • opt (Optimizer) – Built around a minimiser solver, e.g. optx.LBFGS; a least-squares solver raises, its state being untraceable.

  • model (pytree) – Starting model, typically a Scene.

  • probes (mapping of str to callable, optional) – (model) -> array, evaluated at every recorded frame. Traced together under one equinox.filter_jit(), so each must return a JAX array of fixed shape.

  • max_steps (int, optional) – Exact number of solver steps to take, not a ceiling.

  • stride (int, optional) – Record every stride-th step.

  • jit (bool, optional) – Whether to compile the solver step and the probes.

  • callback (callable, optional) – (step, loss, model) -> None, called at every recorded frame.

  • stop_on_diverged (bool, optional) – Whether to stop, warn and return early once the loss goes non-finite.

Returns:

trace.model is the fitted model.

Return type:

FitTrace

Raises:

TypeError – If opt uses a least-squares solver.