{"id":13520939,"url":"https://github.com/patrick-kidger/lineax","last_synced_at":"2025-11-17T15:33:28.965Z","repository":{"id":173188435,"uuid":"638621600","full_name":"patrick-kidger/lineax","owner":"patrick-kidger","description":"Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax","archived":false,"fork":false,"pushed_at":"2025-03-27T17:49:54.000Z","size":239,"stargazers_count":428,"open_issues_count":40,"forks_count":24,"subscribers_count":10,"default_branch":"main","last_synced_at":"2025-04-01T08:33:22.691Z","etag":null,"topics":["equinox","jax","linear-algebra"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/patrick-kidger.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-05-09T18:36:57.000Z","updated_at":"2025-03-31T03:37:05.000Z","dependencies_parsed_at":"2023-09-26T11:34:55.565Z","dependency_job_id":"68e3c567-883d-43aa-a892-c34646569372","html_url":"https://github.com/patrick-kidger/lineax","commit_stats":{"total_commits":99,"total_committers":9,"mean_commits":11.0,"dds":0.3535353535353535,"last_synced_commit":"4d9928378e0ac93a95afdb9bab904db3df312c86"},"previous_names":["google/lineax","patrick-kidger/lineax"],"tags_count":7,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Flineax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Flineax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Flineax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Flineax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/patrick-kidger","download_url":"https://codeload.github.com/patrick-kidger/lineax/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247809964,"owners_count":20999816,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["equinox","jax","linear-algebra"],"created_at":"2024-08-01T06:00:24.752Z","updated_at":"2025-11-17T15:33:23.947Z","avatar_url":"https://github.com/patrick-kidger.png","language":"Python","funding_links":[],"categories":["Models and Projects","Python","Libraries"],"sub_categories":[],"readme":"\u003ch1 align='center'\u003eLineax\u003c/h1\u003e\n\nLineax is a [JAX](https://github.com/google/jax) library for linear solves and linear least squares. That is, Lineax provides routines that solve for $x$ in $Ax = b$. (Even when $A$ may be ill-posed or rectangular.)\n\nFeatures include:\n- PyTree-valued matrices and vectors;\n- General linear operators for Jacobians, transposes, etc.;\n- Efficient linear least squares (e.g. QR solvers);\n- Numerically stable gradients through linear least squares;\n- Support for structured (e.g. symmetric) matrices;\n- Improved compilation times;\n- Improved runtime of some algorithms;\n- Support for both real-valued and complex-valued inputs;\n- All the benefits of working with JAX: autodiff, autoparallelism, GPU/TPU support, etc.\n\n## Installation\n\n```bash\npip install lineax\n```\n\nRequires Python 3.10+, JAX 0.4.38+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.11.10+.\n\n## Documentation\n\nAvailable at [https://docs.kidger.site/lineax](https://docs.kidger.site/lineax).\n\n## Quick examples\n\nLineax can solve a least squares problem with an explicit matrix operator:\n\n```python\nimport jax.random as jr\nimport lineax as lx\n\nmatrix_key, vector_key = jr.split(jr.PRNGKey(0))\nmatrix = jr.normal(matrix_key, (10, 8))\nvector = jr.normal(vector_key, (10,))\noperator = lx.MatrixLinearOperator(matrix)\nsolution = lx.linear_solve(operator, vector, solver=lx.QR())\n```\n\nor Lineax can solve a problem without ever materializing a matrix, as done in this\nquadratic solve:\n\n```python\nimport jax\nimport lineax as lx\n\nkey = jax.random.PRNGKey(0)\ny = jax.random.normal(key, (10,))\n\ndef quadratic_fn(y, args):\n  return jax.numpy.sum((y - 1)**2)\n\ngradient_fn = jax.grad(quadratic_fn)\nhessian = lx.JacobianLinearOperator(gradient_fn, y, tags=lx.positive_semidefinite_tag)\nsolver = lx.CG(rtol=1e-6, atol=1e-6)\nout = lx.linear_solve(hessian, gradient_fn(y, args=None), solver)\nminimum = y - out.value\n```\n\n## Citation\n\nIf you found this library to be useful in academic work, then please cite: ([arXiv link](https://arxiv.org/abs/2311.17283))\n\n```bibtex\n@article{lineax2023,\n    title={Lineax: unified linear solves and linear least-squares in JAX and Equinox},\n    author={Jason Rader and Terry Lyons and Patrick Kidger},\n    journal={\n        AI for science workshop at Neural Information Processing Systems 2023,\n        arXiv:2311.17283\n    },\n    year={2023},\n}\n```\n\n(Also consider starring the project on GitHub.)\n\n## See also: other libraries in the JAX ecosystem\n\n**Always useful**  \n[Equinox](https://github.com/patrick-kidger/equinox): neural networks and everything not already in core JAX!  \n[jaxtyping](https://github.com/patrick-kidger/jaxtyping): type annotations for shape/dtype of arrays.  \n\n**Deep learning**  \n[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  \n[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).  \n[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).  \n[paramax](https://github.com/danielward27/paramax): parameterizations and constraints for PyTrees.  \n\n**Scientific computing**  \n[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  \n[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  \n[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.  \n[sympy2jax](https://github.com/patrick-kidger/sympy2jax): SymPy\u003c-\u003eJAX conversion; train symbolic expressions via gradient descent.  \n[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  \n\n**Awesome JAX**  \n[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  \n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Flineax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpatrick-kidger%2Flineax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Flineax/lists"}