nlls_gram
Levenberg-Marquardt nonlinear least squares for JAX pytrees, aimed at solving systems of equations. The core use case is underdetermined systems — more parameters than residuals — where a zero-residual root is not unique and something must select which interpolant is returned. Two solvers differ in where that selection lives:
RidgeLevenbergMarquardtputs it in the objective, minimizing (|r(x)|^2 + \lambda,|x_m|_W^2) for a positive-definite metric (W) on the metric block of (x = [x_m; x_f]) (the free block stays unpenalized). Annealing (\lambda) toward zero converges to the minimum-seminorm — e.g. minimum-RKHS-norm — interpolant, by classical nonlinear Tikhonov regularization. Every inner problem is a well-posed NLLS.LevenbergMarquardtputs it in the damping geometry: standard damped LM whose trust region is measured in (W), so the small-damping Gauss-Newton limit is the minimum-(W)-norm correction.
Both take a residual over (x), (x, args), or (x, args, p), flatten any
pytree x, expose per-step update(...) and an internally jitted solve(...)
loop with callbacks and multi-start, and differentiate solve(...).x with
respect to p through a custom implicit rule — no unrolling.
Install
uv add nlls-gram
For GPU use, install the JAX accelerator build that matches your hardware:
uv add nlls-gram "jax[cuda13]"
Minimum-RKHS-norm interpolation
import jax.numpy as jnp
from nlls_gram import AnnealRidge, RidgeLevenbergMarquardt, RepeatedFactorMetric
# W = blockdiag(K, K): the RKHS seminorm over two coefficient blocks. The
# constructor takes the FACTOR; shift a semidefinite K by epsilon*I first.
metric = RepeatedFactorMetric(jnp.linalg.cholesky(K, upper=True), repeats=2)
solver = RidgeLevenbergMarquardt(collocation_residual, metric=metric, ridge=1e-4)
# Anneal the ridge toward the interpolating limit on stationarity.
anneal = AnnealRidge(ridge_floor=1e-10)
result = solver.solve(x0, callback=anneal, user_state=anneal.init_state(),
gtol=1e-8, atol=1e-8)
General nonlinear least squares
import jax, jax.numpy as jnp
from nlls_gram import LevenbergMarquardt
def residual(x, args, p):
return args["design"] @ x - p["target"]
solver = LevenbergMarquardt(residual)
result = solver.solve(jnp.zeros(8), {"design": design}, p={"target": y},
max_steps=200, atol=1e-8)
# The solution is differentiable in p, at a cost independent of max_steps.
sensitivity = jax.grad(
lambda p: jnp.sum(solver.solve(jnp.zeros(8), {"design": design}, p=p,
max_steps=200, atol=1e-8).x ** 2)
)({"target": y})
Pass metric= to weight the damping geometry.
Linear solvers
The same typed configs serve both solvers, forward and in the implicit-AD
role: Cholesky() (the default, auto-selecting the smaller of the dual and
normal systems), QR() (damping-row QR — stable at tiny damping and
rank-safe), CG(precond) and GramCG(precond) (matrix-free in parameter and
residual space), and SVD() for rank-deficient tangents. A knob that exists
for only one method is a field on that method, so it cannot be passed with
another.
Docs
https://highdimensionaleconlab.github.io/nlls_gram/
License
MIT.
1 maintainer
Metadata
1 maintainer
Release files for nlls-gram 2.9.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| nlls_gram-2.9.0.tar.gz | 230.9 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| nlls_gram-2.9.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 308.8 kB
Release files / nlls_gram-2.9.0.tar.gz
| Download URL | nlls_gram-2.9.0.tar.gz |
|---|---|
| Size | 230.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
821a15d4025ffd9daf40f413fd593c568f129c5ccd24efde8b0cce9544610182
|
|
BLAKE2b-256 checksum How to use checksums |
2c44e7409ddd699fcf0b0d00f0fd04dc8d387e00213acd23d888f14ea136e4b0
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Aug 27, 2026.
Transparency logRelease files / nlls_gram-2.9.0-py3-none-any.whl
| Download URL | nlls_gram-2.9.0-py3-none-any.whl |
|---|---|
| Size | 77.9 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
5ec3747e721fc4390be2e206470821c7c3cb1154a8aa29d71d336ee677e9eab2
|
|
BLAKE2b-256 checksum How to use checksums |
01f57a0ae75f6ebe3ab5b83f9e07a350c7ce81723230b54a62b0ccb5cb152cb7
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Aug 27, 2026.
Transparency log
