jaxls

jaxls#

Pyright status

jaxls is a solver for sparse, constrained, and/or non-Euclidean least squares problems in JAX.

Please install from source:

# Core package only
pip install "git+https://github.com/brentyi/jaxls.git"

# Or, with dev dependencies
pip install "git+https://github.com/brentyi/jaxls.git#egg=jaxls[dev,docs]"

We recommend Python >=3.12. The test matrix also covers Python 3.10 and 3.11, but this uses a transpiled version of the source code that may be deprecated in the future.

Goals#

jaxls is a research artifact and intentionally lightweight. It’s designed to be:

  • Hackable but fast. Performance is enabled by analyzing and exploiting problem structure: jaxls automatically vectorizes repeated cost and variable operations, while translating sparse cost/variable relationships into sparse matrix operations.

  • Python-native. jaxls combines a functional, PyTree-first implementation with recent Python typing constructs. Its API is type-safe, compatible with standard JAX function transforms, and more concise than traditional optimization tools.

Features#

We currently support:

  • Automatic sparse Jacobians, defined analytically or via autodiff. See Custom Jacobians and Sparse matrices.

  • Optimization on manifolds, including SO(2), SO(3), SE(2), SE(3). See Non-Euclidean variables.

  • Nonlinear solvers: Levenberg-Marquardt and Gauss-Newton.

  • Automatic Schur-complement variable elimination for problems with a dominant block-diagonal type, such as landmarks in bundle adjustment. See Schur complement elimination.

  • Linear subproblem solvers. See Tips and gotchas for selection guidance.

    • Sparse iterative with Conjugate Gradient (recommended for most problems).

      • Block and point Jacobi preconditioning; inexact Newton via Eisenstat-Walker.

    • Dense Cholesky (fast for small problems).

    • Sparse Cholesky on CPU (CHOLMOD).

  • Augmented Lagrangian solver for constrained problems. See Constraints.

    • Automatic adaptive penalties for equalities and inequalities: h(x) = 0, g(x) <= 0.

Acknowledgements#

Thanks to Nick Heppert, Chung Min Kim, and Yuqing Du for technical feedback.

jaxls is inspired by libraries like GTSAM, Ceres Solver, minisam, SwiftFusion, and g2o.

Algorithmic references:

Contents