From ef7a61436dce8cdf3db5f288803bbe59d7e3ff4f Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Thu, 27 Aug 2026 13:06:28 -0700 Subject: [PATCH] feat: DrawNNXModule takes a wrt filter for the drawn state DrawNNXModule(module_cls, *args, wrt=None, **kwargs): wrt is a keyword-only nnx filter selecting the state drawn from each redrawn module, defaulting to nnx.Param (None resolves at draw time so the flax import stays lazy). Pass the same filter used to split the solver's x0 when some Params are frozen into the nondiff state -- e.g. nnx.All(nnx.Param, nnx.Not(nnx.PathContains("g"))) -- so drawn states match x0's pytree structure. wrt joins the value-hash spec alongside (module_cls, args, kwargs). Version 2.9.0. Co-Authored-By: Mecha Perla (Claude) Claude-Session: https://claude.ai/code/session_01JqaCzfRpewquuigfnNkP4n --- pyproject.toml | 2 +- src/nlls_gram/multi_start.py | 43 +++++++++++++++++++++++------------- tests/test_nnx_gram_lm.py | 18 +++++++++++++++ uv.lock | 2 +- 4 files changed, 48 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b0869b1..0aac211 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "nlls-gram" -version = "2.8.0" +version = "2.9.0" description = "Metric-aware Levenberg-Marquardt nonlinear least-squares for JAX" readme = "README.md" license = "MIT" diff --git a/src/nlls_gram/multi_start.py b/src/nlls_gram/multi_start.py index 0da9263..56add52 100644 --- a/src/nlls_gram/multi_start.py +++ b/src/nlls_gram/multi_start.py @@ -91,49 +91,62 @@ 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. Non-``Param`` - Variables (scaling constants, statistics) are excluded from the drawn state -- - the residual's ``nnx.merge`` supplies them alongside the graphdef. Use it instead - of hand-rolling a re-init closure per driver:: + ``module_cls(*args, rngs=nnx.Rngs(key), **kwargs)`` and returns its ``wrt`` + state (default ``nnx.Param``) as the new solver start, passing ``args`` + through unchanged. Non-``Param`` Variables (scaling constants, statistics) + are excluded from the drawn state -- the residual's ``nnx.merge`` supplies + them alongside the graphdef. 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) + ``wrt`` is a keyword-only ``nnx`` filter (not forwarded to ``module_cls``): + when the solver's ``x0`` was split with a filter narrower than ``nnx.Param`` + (e.g. ``nnx.All(nnx.Param, nnx.Not(nnx.PathContains("g")))`` to freeze a + Param into the nondiff state), pass the same filter so drawn states match + ``x0``'s pytree structure. + 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). + Value-hashable on ``(module_cls, wrt, 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): + def __init__(self, module_cls, *args, wrt=None, **kwargs): self.module_cls = module_cls self.args = args + # None means nnx.Param, resolved at draw time to keep the flax import lazy. + self.wrt = wrt 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, ...) + wrt = nnx.Param if self.wrt is None else self.wrt + _, theta, _ = nnx.split(module, wrt, ...) return theta, args_old def __hash__(self): - return hash((self.module_cls, _typed_key(self.args), _typed_key(self.kwargs))) + return hash( + (self.module_cls, self.wrt, _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 self.wrt == other.wrt and _typed_key(self.args) == _typed_key(other.args) and _typed_key(self.kwargs) == _typed_key(other.kwargs) ) diff --git a/tests/test_nnx_gram_lm.py b/tests/test_nnx_gram_lm.py index 11d1b33..c8cd92e 100644 --- a/tests/test_nnx_gram_lm.py +++ b/tests/test_nnx_gram_lm.py @@ -273,6 +273,24 @@ def residual(theta, args, p): assert bool(result.multi_start.accepted) +def test_draw_nnx_module_wrt_filter(): + wrt = nnx.All(nnx.Param, nnx.Not(nnx.PathContains("head"))) + _, theta_0, _ = nnx.split(CurveMLP(rngs=nnx.Rngs(1)), wrt, ...) + + draw = DrawNNXModule(CurveMLP, wrt=wrt) + theta, args_out = draw(jax.random.key(7), None, ("args",)) + assert args_out == ("args",) + # Only the wrt-selected Params are drawn; the rest rides in the frozen state. + assert jax.tree_util.tree_structure(theta) == jax.tree_util.tree_structure(theta_0) + + # wrt joins the value-hash spec. + assert DrawNNXModule(CurveMLP, wrt=wrt) == DrawNNXModule(CurveMLP, wrt=wrt) + assert hash(DrawNNXModule(CurveMLP, wrt=wrt)) == hash( + DrawNNXModule(CurveMLP, wrt=wrt) + ) + assert DrawNNXModule(CurveMLP, wrt=wrt) != DrawNNXModule(CurveMLP) + + @pytest.mark.parametrize("parallel", [False, True]) def test_multi_start_draw_nnx_module_recovers_from_bad_init(parallel): ts = jnp.linspace(-1.0, 1.0, 32) diff --git a/uv.lock b/uv.lock index 0fd1bbb..1eee888 100644 --- a/uv.lock +++ b/uv.lock @@ -726,7 +726,7 @@ wheels = [ [[package]] name = "nlls-gram" -version = "2.8.0" +version = "2.9.0" source = { editable = "." } dependencies = [ { name = "jax" },