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" },