diff --git a/pyproject.toml b/pyproject.toml index 868a701..b0869b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "nlls-gram" -version = "2.7.0" +version = "2.8.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 d420473..0da9263 100644 --- a/src/nlls_gram/multi_start.py +++ b/src/nlls_gram/multi_start.py @@ -92,7 +92,9 @@ class DrawNNXModule: 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 + 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:: draw = DrawNNXModule(SequentialMLP, settings, dtype=dtype) @@ -122,7 +124,7 @@ 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) + _, theta, _ = nnx.split(module, nnx.Param, ...) return theta, args_old def __hash__(self): diff --git a/tests/test_nnx_gram_lm.py b/tests/test_nnx_gram_lm.py index b0023c9..11d1b33 100644 --- a/tests/test_nnx_gram_lm.py +++ b/tests/test_nnx_gram_lm.py @@ -236,6 +236,43 @@ def run(draw): assert jnp.allclose(got, want) +class ScaledCurveMLP(nnx.Module): + def __init__(self, *, scale=2.0, rngs: nnx.Rngs): + self.scale = nnx.Variable(jnp.asarray(scale)) + self.hidden = nnx.Linear(1, 8, rngs=rngs) + self.head = nnx.Linear(8, 1, rngs=rngs) + + def __call__(self, x): + return self.scale[...] * self.head(nnx.tanh(self.hidden(x[:, None])))[:, 0] + + +def test_draw_nnx_module_excludes_non_param_variables(): + graphdef, theta_0, nondiff = nnx.split( + ScaledCurveMLP(rngs=nnx.Rngs(1)), nnx.Param, ... + ) + assert len(jax.tree.leaves(nondiff)) == 1 + + draw = DrawNNXModule(ScaledCurveMLP) + theta, args_out = draw(jax.random.key(7), None, ("args",)) + assert args_out == ("args",) + # Only the Param leaves are drawn; the residual's merge supplies nondiff. + assert jax.tree_util.tree_structure(theta) == jax.tree_util.tree_structure(theta_0) + + ts = jnp.linspace(-1.0, 1.0, 32) + ys = jnp.sin(2.0 * ts) + theta_bad = jax.tree.map(lambda leaf: leaf * jnp.nan, theta_0) + + def residual(theta, args, p): + ts, ys = args + return nnx.merge(graphdef, theta, nondiff)(ts) - ys + + solver = LevenbergMarquardt(residual, init_damping=1e-2) + ms = MultiStart(key=jax.random.key(2), num_starts=4, draw=draw) + result = solver.solve(theta_bad, (ts, ys), max_steps=200, atol=5e-3, multi_start=ms) + assert int(result.status) == LMStatus.CONVERGED + assert bool(result.multi_start.accepted) + + @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 2aff017..0fd1bbb 100644 --- a/uv.lock +++ b/uv.lock @@ -726,7 +726,7 @@ wheels = [ [[package]] name = "nlls-gram" -version = "2.7.0" +version = "2.8.0" source = { editable = "." } dependencies = [ { name = "jax" },