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:

RidgeLevenbergMarquardt puts 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.
LevenbergMarquardt puts 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.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

nlls_gram-2.8.0.tar.gz (230.5 kB view details)

Uploaded Aug 18, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

nlls_gram-2.8.0-py3-none-any.whl (77.6 kB view details)

Uploaded Aug 18, 2026 Python 3

File details

Details for the file nlls_gram-2.8.0.tar.gz.

File metadata

Download URL: nlls_gram-2.8.0.tar.gz
Upload date: Aug 18, 2026
Size: 230.5 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for nlls_gram-2.8.0.tar.gz
Algorithm	Hash digest
SHA256	`b2cd7776c7b6b58a2ee0ad3d20ddecd809de70db45625e042600b20f8b73b76d`
MD5	`a503bb96a5f740ee8b9f0b37efc9ed45`
BLAKE2b-256	`f0401041bbcc2cd32da39e75cf18797a7f57bd2f198c21608a2e35856a6a36dd`

See more details on using hashes here.

Provenance

The following attestation bundles were made for nlls_gram-2.8.0.tar.gz:

Publisher: publish.yml on HighDimensionalEconLab/nlls_gram

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: nlls_gram-2.8.0.tar.gz
- Subject digest: b2cd7776c7b6b58a2ee0ad3d20ddecd809de70db45625e042600b20f8b73b76d
- Sigstore transparency entry: 2507745775
- Sigstore integration time: Aug 18, 2026
Source repository:
- Permalink: HighDimensionalEconLab/nlls_gram@41fb8863448313e50929482086aafdaeed2bb4e6
- Branch / Tag: refs/tags/v2.8.0
- Owner: https://github.com/HighDimensionalEconLab
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@41fb8863448313e50929482086aafdaeed2bb4e6
- Trigger Event: release