jnlr 0.1.1

Jax-based library for nonlinear reconciliation and learning

JNLR Jax-based non-linear reconciliation and learning

J-NLR is a Python library for non-linear reconciliation, learning, and geometric analysis on constraint manifolds. Built on JAX, it leverages automatic differentiation and GPU/TPU acceleration to efficiently project predicted values onto surfaces defined by implicit constraints $f(z) = 0$.

📚 Documentation - Full API reference, examples, and interactive notebooks

Key Features

• Non-linear Reconciliation: Multiple solvers (Augmented Lagrangian, curvature-aware Newton, vanilla projections) for projecting forecasts onto constraint manifolds
• SHOULD Analysis: Curvature-based methods to determine when reconciliation is beneficial—verify if RMSE is guaranteed to reduce before applying corrections
• Manifold Sampling: Sample from explicit (graph) or implicit manifolds using volume-weighted sampling, Latin hypercube, or Langevin dynamics on the constraint surface
• Mesh Generation: Create triangulated meshes from explicit parameterizations for visualization and geodesic computation
• Geodesics: Compute geodesic distances and shortest paths on manifolds via exact MMP algorithm or fast graph-based approximations; includes probabilistic scores like pointcloud geodesic distance
• Visualization: Interactive 3D rendering of manifolds, projections, and geodesic paths with Plotly
• JAX-native: Fully JIT-compiled and vectorized (vmap) for high-performance batch processing

Running the notebooks

If you want to run the notebooks, the suggested way is to install uv packet manager, cloning the repo and, from a terminal: uv pip install -e .

Citation

If you use JNLR in academic work, please cite the associated paper:

Lorenzo Nespoli, Anubhab Biswas, Roberto Rocchetta, and Vasco Medici.
"Nonlinear reconciliation: Error reduction theorems."
Transactions on Machine Learning Research (TMLR), 2026.
OpenReview: https://openreview.net/forum?id=dXRWuogm3J

BibTeX

@article{nespoli2026nonlinear_reconciliation,
  title   = {Nonlinear reconciliation: Error reduction theorems},
  author  = {Nespoli, Lorenzo and Biswas, Anubhab and Rocchetta, Roberto and Medici, Vasco},
  journal = {Transactions on Machine Learning Research},
  year    = {2026},
  url     = {https://openreview.net/forum?id=dXRWuogm3J},
  note    = {Accepted by TMLR}
}

Acknowledgements

This work has been funded by the Swiss State Secretariat for Education, Research and Innovation (SERI) under the Swiss contribution to the Horizon Europe projects DR-RISE (Horizon Europe, Grant Agreement No. 101104154) and REEFLEX (Horizon Europe, Grant Agreement No. 101096192).