Multi-Start

solve(multi_start=MultiStart(...)) retries a failed solve from fresh initial conditions, or races several starts in parallel and keeps the best one. Either way it returns a single LMSolveResult — the selected solution — with diagnostics attached at result.multi_start (a MultiStartInfo). With multi_start=None (the default) nothing changes: the solve takes exactly the plain code path.

import jax
import jax.numpy as jnp
from nlls_gram import MultiStart, UnderdeterminedLevenbergMarquardt

def residual(theta, args, p):
    return theta[0] ** 2 - p  # stalls when started at theta = 0

def draw(key, x, args):
    # Fresh initial condition; args may be redrawn too (see below).
    return jax.random.uniform(key, x.shape, x.dtype, 0.5, 3.0), args

solver = UnderdeterminedLevenbergMarquardt(residual)
ms = MultiStart(key=jax.random.key(0), num_starts=5, draw=draw)
result = solver.solve(jnp.zeros(1), p=jnp.asarray(4.0), atol=1e-8, multi_start=ms)
result.multi_start.attempt        # which start won (0 = your x0)
result.multi_start.attempts_run   # how many solves actually ran

draw and accept are jit static arguments, so they key the compile by their __hash__/__eq__. Plain functions, lambdas, and closures hash by identity — define them once at setup scope, never as fresh lambdas at the call site, or each one recompiles. A value-hashable hook keys by value instead: DrawNNXModule (below) shares one compile across equal specs. A MultiStart object itself is cheap to rebuild per call (only the hooks' cache keys matter), and its key is ordinary traced data — new keys, new x0/args values, and (in sequential mode) a different num_starts among values above one all reuse the compiled solve. Crossing num_starts = 1 to N > 1 (or back) compiles once more: the single-start form never draws, so it is a structurally different program.

The draw contract

def draw(key, x_old, args_old):
    ...
    return x_new, args_new

draw must be traceable and type-stable: the returned (x, args) must match the input pytree structure, shapes, and dtypes exactly (checked up front with an abstract jax.eval_shape trace under jit=True; the first concrete draw is checked under jit=False). What it draws is up to you — reinitialize a network, perturb the previous start, resample the data inside args, or any combination. p is deliberately not an input and cannot change across starts: it is the differentiation target, and the implicit gradient is taken at a single fixed p.

In sequential mode draw receives the previous attempt's initial values — the original (x0, args) for the first retry, then each drawn (x, args) in turn — never the solver-mutated result.x/result.args. In parallel mode every lane draws from the original (x0, args).

A flax nnx reinitialization draw:

def draw(key, x_old, args_old):
    _, theta = nnx.split(PolicyMLP(settings, rngs=nnx.Rngs(key)), nnx.Param)
    return theta, args_old

DrawNNXModule packages exactly this draw so you skip the per-driver closure:

from nlls_gram import DrawNNXModule

draw = DrawNNXModule(PolicyMLP, settings, dtype=dtype)  # equal specs share one compile

It rebuilds module_cls(*args, rngs=nnx.Rngs(key), **kwargs) on each retry and returns its nnx.Param state, passing args through unchanged. The drawn state must be type-stable against x0 (same structure, shapes, dtypes), so construct the module with a matching param_dtype/dtype (e.g. thread dtype= through). Unlike a fresh closure it is value-hashable on (module_cls, args, kwargs), so equal specs share a single jit compilation instead of recompiling.

A data-resampling draw (mv2020 style), threading a key inside args:

def draw(key, x_old, args_old):
    init_key, exo_key, carry_key = jax.random.split(key, 3)
    _, theta = nnx.split(PolicyMLP(settings, rngs=nnx.Rngs(init_key)), nnx.Param)
    args_new = args_old.replace(
        exo=simulate_markov_chain(exo_key, s_0, P_cumsum, train_T),
        key=carry_key,
        epoch=jnp.asarray(0, jnp.int32),
    )
    return theta, args_new

The accept hook

By default an attempt succeeds when result.status == LMStatus.CONVERGED. Pass accept to override the test — for example to require a fresh-data validation metric rather than trusting the training tolerance:

def accept(key, result):
    policy = nnx.merge(graphdef, result.x)
    test_exo = simulate_markov_chain(key, s_0_test, P_cumsum, test_T)
    return euler_mean_abs(policy, test_exo) < 5e-5

ms = MultiStart(key=key, num_starts=5, draw=draw, accept=accept)

accept receives its own key (see the schedule below) and must return a scalar boolean-like value; the solver canonicalizes the dtype. Inside accept, result.multi_start is still None (the diagnostics are attached after selection). Note result.info.loss can be stale when a callback replaced x/args after the last update — recompute anything you need at (result.x, result.args).

An accepted-but-nonfinite result never wins: effective success is accept(...) AND isfinite(loss) in both modes.

Sequential vs parallel

parallel changes the selection semantics, not just the execution strategy:

	sequential (default)	parallel=True
execution	one solve at a time, stops at the first success	all num_starts lanes under one vmap
winner	the first accepted attempt	the accepted lane with the lowest loss
all fail	lowest finite loss across attempts	lowest finite loss across lanes
none finite	the last attempt	lane 0
cost	pays only for the attempts run	always pays for num_starts solves (but they run batched)
num_starts	traced — changes among values > 1 never retrace	static — changing it recompiles

Parallel lanes share one vmapped while_loop, so every lane steps until the slowest lane stops: the compiled cost is num_starts x slowest lane, not the sum of each lane's own step count. Budget max_steps accordingly.

The ranking loss is the sum of squared residuals at the returned solution (result.info.loss, or recomputed at (result.x, result.args, p) when a callback is present), masked to +inf when nonfinite; ties break to the lowest attempt index. MultiStartInfo.loss records the winner's value.

Identical draw keys are used in both modes (below), so a draw that ignores (x_old, args_old) produces the same candidate starts sequentially and in parallel; the modes still may pick different winners (first-accepted vs best-of-batch).

Key schedule

For attempt/lane k:

draw_key, accept_key = jax.random.split(jax.random.fold_in(key, k))

Attempt 0 is always the caller's (x0, args) and never consumes its draw_key. The schedule is a documented contract (pinned by tests), so runs are reproducible and an attempt's draws do not depend on how many attempts ran before it.

Differentiation

Gradients with respect to p flow through the selected solution only, via the same implicit rule as a plain solve (see implicit differentiation): the residual is relinearized at the returned (x, args, p), and everything else — the key, the initial conditions, the losing attempts, the diagnostics — has zero tangents.

Because selection is discrete, the derivative is piecewise: it is the chosen basin's implicit derivative, and it jumps when a different start wins (ties, basin switches). The acceptance test and the argmin are not differentiated. Derivatives are only meaningful when the winner actually solved the problem — after an all-fail fallback the rule linearizes at a non-solution.

Interactions

• lm_state warm starts — the caller's lm_state applies to attempt 0; drawn attempts inherit its damping and hyperparameters but the Jacobian cache is invalidated (it described a different (x, args)). In parallel mode the cache is dropped on all lanes: under vmap the cache-reuse branch is a select that evaluates both sides, so a warm cache cannot save work.
• save_steps — composes; the winner's (unbatched) histories are returned. Parallel mode materializes num_starts history buffers during the solve, and sequential mode briefly holds two.
• Outer vmap — sequential multi-start composes with an outer vmap (e.g. one multi-start solve per sample); as with any vmapped while_loop, all lanes wait for the slowest sample's schedule. Parallel-inside-vmap nests two batch axes — fine for small num_starts x batch, memory-hungry beyond that.
• jit=False — both modes run eagerly for debugging; sequential mode then calls draw lazily, only when a retry actually happens.

API

nlls_gram.MultiStart dataclass

Multi-start configuration for solve(multi_start=...).

draw(key, x, args) -> (x_new, args_new) generates a fresh initial condition; it must be traceable and type-stable (returning the same pytree structure, shapes, and dtypes as its (x, args) inputs). accept(key, result) -> bool optionally overrides the success test (default: result.status == LMStatus.CONVERGED); it receives its own key so it can draw fresh validation data, and may return any scalar boolean-like value. Sequential mode (parallel=False) solves from (x0, args) and retries on failure, chaining each attempt's initial values into the next draw; parallel mode solves all num_starts lanes under vmap (lane 0 = the caller's (x0, args), the rest drawn from the originals) and selects the accepted lane with the lowest loss. The key schedule is draw_key, accept_key = jax.random.split(jax.random.fold_in(key, k)) for attempt k.

draw and accept enter the jit cache by identity (like callback): define them once at setup scope, not inline per call. MultiStart is not a pytree -- solve unpacks it before tracing, with key the only traced field.

Source code in src/nlls_gram/gram_lm.py

@dataclass(frozen=True, eq=False)
class MultiStart:
    """Multi-start configuration for ``solve(multi_start=...)``.

    ``draw(key, x, args) -> (x_new, args_new)`` generates a fresh initial
    condition; it must be traceable and type-stable (returning the same pytree
    structure, shapes, and dtypes as its ``(x, args)`` inputs). ``accept(key,
    result) -> bool`` optionally overrides the success test (default:
    ``result.status == LMStatus.CONVERGED``); it receives its own key so it can
    draw fresh validation data, and may return any scalar boolean-like value.
    Sequential mode (``parallel=False``) solves from ``(x0, args)`` and retries
    on failure, chaining each attempt's *initial* values into the next
    ``draw``; parallel mode solves all ``num_starts`` lanes under ``vmap``
    (lane 0 = the caller's ``(x0, args)``, the rest drawn from the originals)
    and selects the accepted lane with the lowest loss. The key schedule is
    ``draw_key, accept_key = jax.random.split(jax.random.fold_in(key, k))``
    for attempt ``k``.

    ``draw`` and ``accept`` enter the jit cache by identity (like
    ``callback``): define them once at setup scope, not inline per call.
    ``MultiStart`` is not a pytree -- ``solve`` unpacks it before tracing, with
    ``key`` the only traced field.
    """

    key: Any
    num_starts: int
    draw: Any = None
    accept: Any = None
    parallel: bool = False

    def __post_init__(self):
        if isinstance(self.num_starts, bool) or not isinstance(self.num_starts, int):
            raise ValueError("num_starts must be a Python int >= 1")
        if self.num_starts < 1:
            raise ValueError("num_starts must be a Python int >= 1")
        if self.num_starts > 1 and self.draw is None:
            raise ValueError(
                "num_starts > 1 requires draw; pass "
                "draw=(key, x, args) -> (x_new, args_new)"
            )
        if self.draw is not None and not callable(self.draw):
            raise TypeError("draw must be callable")
        if self.accept is not None and not callable(self.accept):
            raise TypeError("accept must be callable")

nlls_gram.MultiStartInfo dataclass

Diagnostics attached to LMSolveResult.multi_start by a multi-start solve.

attempt is the winning attempt/lane index (0 = the caller's (x0, args)), accepted whether the winner passed the success test (MultiStart.accept, or status == LMStatus.CONVERGED), and attempts_run how many starts were solved (sequential mode stops at the first success; parallel mode always runs num_starts). loss is the ranking loss selection used: the sum of squared residuals at the returned solution, masked to +inf when nonfinite. Note accepted describes the multi-start success test, not LMInfo.accepted (last-step acceptance).

Source code in src/nlls_gram/gram_lm.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class MultiStartInfo:
    """Diagnostics attached to ``LMSolveResult.multi_start`` by a multi-start solve.

    ``attempt`` is the winning attempt/lane index (0 = the caller's
    ``(x0, args)``), ``accepted`` whether the winner passed the success test
    (``MultiStart.accept``, or ``status == LMStatus.CONVERGED``), and
    ``attempts_run`` how many starts were solved (sequential mode stops at the
    first success; parallel mode always runs ``num_starts``). ``loss`` is the
    ranking loss selection used: the sum of squared residuals at the returned
    solution, masked to ``+inf`` when nonfinite. Note ``accepted`` describes
    the multi-start success test, not ``LMInfo.accepted`` (last-step
    acceptance).
    """

    attempt: jax.Array
    accepted: jax.Array
    attempts_run: jax.Array
    loss: jax.Array

nlls_gram.DrawNNXModule

Multi-start draw hook re-initializing a flax nnx.Module from a fresh key.

Given a MultiStart retry key, builds module_cls(*args, rngs=nnx.Rngs(key), **kwargs) and returns its nnx.Param state as the new solver start, passing args through unchanged. Use it instead of hand-rolling a re-init closure per driver::

draw = DrawNNXModule(SequentialMLP, settings, dtype=dtype)
ms = MultiStart(key=key, num_starts=5, draw=draw)

The drawn parameter state must be type-stable against the solver's x0 (same pytree structure, shapes, and dtypes) -- construct the module with a matching param_dtype/dtype (e.g. pass dtype= through). The paired nnx.GraphDef used by the residual's nnx.merge must come from the same module_cls(*args, **kwargs) spec.

Value-hashable on (module_cls, args, kwargs) with jit's strict-type semantics (1, 1.0, and True key distinct compilations): equal specs compare equal and share one jit compilation instead of recompiling per instance (a fresh closure would not). args/kwargs must be hashable for that sharing, and their values must not be mutated after construction (a stale key would reuse the wrong compile); unhashable specs still work but recompile per instance. Requires flax installed (imported lazily on first draw).

Source code in src/nlls_gram/gram_lm.py

class DrawNNXModule:
    """Multi-start ``draw`` hook re-initializing a flax ``nnx.Module`` from a fresh key.

    Given a ``MultiStart`` retry key, builds
    ``module_cls(*args, rngs=nnx.Rngs(key), **kwargs)`` and returns its ``nnx.Param``
    state as the new solver start, passing ``args`` through unchanged. Use it instead
    of hand-rolling a re-init closure per driver::

        draw = DrawNNXModule(SequentialMLP, settings, dtype=dtype)
        ms = MultiStart(key=key, num_starts=5, draw=draw)

    The drawn parameter state must be type-stable against the solver's ``x0`` (same
    pytree structure, shapes, and dtypes) -- construct the module with a matching
    ``param_dtype``/``dtype`` (e.g. pass ``dtype=`` through). The paired
    ``nnx.GraphDef`` used by the residual's ``nnx.merge`` must come from the same
    ``module_cls(*args, **kwargs)`` spec.

    Value-hashable on ``(module_cls, args, kwargs)`` with jit's strict-type semantics
    (``1``, ``1.0``, and ``True`` key distinct compilations): equal specs compare equal
    and share one jit compilation instead of recompiling per instance (a fresh closure
    would not). ``args``/``kwargs`` must be hashable for that sharing, and their values
    must not be mutated after construction (a stale key would reuse the wrong compile);
    unhashable specs still work but recompile per instance. Requires ``flax`` installed
    (imported lazily on first draw).
    """

    def __init__(self, module_cls, *args, **kwargs):
        self.module_cls = module_cls
        self.args = args
        self.kwargs = tuple(sorted(kwargs.items()))

    def __call__(self, key, x_old, args_old):
        from flax import nnx

        module = self.module_cls(*self.args, rngs=nnx.Rngs(key), **dict(self.kwargs))
        _, theta = nnx.split(module, nnx.Param)
        return theta, args_old

    def __hash__(self):
        return hash((self.module_cls, _typed_key(self.args), _typed_key(self.kwargs)))

    def __eq__(self, other):
        return (
            isinstance(other, DrawNNXModule)
            and self.module_cls is other.module_cls
            and _typed_key(self.args) == _typed_key(other.args)
            and _typed_key(self.kwargs) == _typed_key(other.kwargs)
        )