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.
Row-wise residuals
For an expensive per-state residual vmapped over a simulated path, replace
that jax.vmap with vmap_rows: the same values, with each row's gradient
computed once per linearization and reused across the Jacobian assembly
instead of one pass of the row function per row and Jacobian column (or
cotangent row).
from nlls_gram import vmap_rows
@vmap_rows(in_axes=(0, None, None))
def euler_residual(state, policy, p): ...
Docs
https://highdimensionaleconlab.github.io/nlls_gram/
License
MIT.
1 maintainer
Metadata
1 maintainer
Release files for nlls-gram 2.12.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | |
|---|---|---|---|---|
| nlls_gram-2.12.1.tar.gz | 254.7 kB | With provenance | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| nlls_gram-2.12.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 336.1 kB
Release files / nlls_gram-2.12.1.tar.gz
| Download URL | nlls_gram-2.12.1.tar.gz |
|---|---|
| Size | 254.7 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
d4b07e3bded314d81bf1d2185c34f7132b9af56f8f93ba0992ce9b0654d5d332
|
|
BLAKE2b-256 checksum How to use checksums |
6b206f943e9d4c53335b4e10e8911e6ccc0d668ad34a38e998060fa68007c341
|
| 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 Oct 19, 2026.
Transparency logRelease files / nlls_gram-2.12.1-py3-none-any.whl
| Download URL | nlls_gram-2.12.1-py3-none-any.whl |
|---|---|
| Size | 81.4 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
a272bae908980f35ff17280a430b8a87d06cd25267e04f3d49f3183c088636a6
|
|
BLAKE2b-256 checksum How to use checksums |
6ed9c40b6787a00d42f27d5a64404533c620446e4d4240c7d5e0ca5fb6fac6e3
|
| 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 Oct 19, 2026.
Transparency logSecurity for nlls-gram 2.12.1
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.
Full provenance
All 2 files in this release have been uploaded with provenance metadata.
Attestation for nlls_gram-2.12.1.tar.gz
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 Oct 19, 2026.
Transparency logA note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
