{"id":13563812,"url":"https://github.com/patrick-kidger/equinox","last_synced_at":"2026-01-25T02:14:31.415Z","repository":{"id":37230759,"uuid":"390573152","full_name":"patrick-kidger/equinox","owner":"patrick-kidger","description":"Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/","archived":false,"fork":false,"pushed_at":"2025-05-08T20:04:33.000Z","size":41746,"stargazers_count":2345,"open_issues_count":192,"forks_count":161,"subscribers_count":20,"default_branch":"main","last_synced_at":"2025-05-08T21:22:50.947Z","etag":null,"topics":["deep-learning","equinox","jax","neural-networks"],"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":".github/FUNDING.yml","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,"zenodo":null},"funding":{"github":["patrick-kidger"]}},"created_at":"2021-07-29T02:21:39.000Z","updated_at":"2025-05-08T20:05:48.000Z","dependencies_parsed_at":"2023-12-30T03:28:58.805Z","dependency_job_id":"9ad4a190-7198-453a-b398-ba27f8aa931f","html_url":"https://github.com/patrick-kidger/equinox","commit_stats":{"total_commits":457,"total_committers":27,"mean_commits":"16.925925925925927","dds":"0.25820568927789933","last_synced_commit":"83a1aacb4b25f5bbf6faa81c9bc7e2ef2b76f700"},"previous_names":[],"tags_count":49,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fequinox","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fequinox/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fequinox/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fequinox/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/patrick-kidger","download_url":"https://codeload.github.com/patrick-kidger/equinox/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":253745197,"owners_count":21957319,"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":["deep-learning","equinox","jax","neural-networks"],"created_at":"2024-08-01T13:01:23.529Z","updated_at":"2025-11-17T15:31:57.460Z","avatar_url":"https://github.com/patrick-kidger.png","language":"Python","funding_links":["https://github.com/sponsors/patrick-kidger"],"categories":["Python","Machine Learning Framework","其他_机器学习与深度学习","Libraries","1. Core Frameworks \u0026 Libraries"],"sub_categories":["General Purpose Framework"],"readme":"\u003ch1 align='center'\u003eEquinox\u003c/h1\u003e\n\nEquinox is your one-stop [JAX](https://github.com/google/jax) library, for everything you need that isn't already in core JAX:\n\n- neural networks (or more generally any model), with easy-to-use PyTorch-like syntax;\n- filtered APIs for transformations;\n- useful PyTree manipulation routines;\n- advanced features like runtime errors;\n\nand best of all, Equinox isn't a framework: everything you write in Equinox is compatible with anything else in JAX or the ecosystem.\n\nIf you're completely new to JAX, then start with this [CNN on MNIST example](https://docs.kidger.site/equinox/examples/mnist/).\n\n_Coming from [Flax](https://github.com/google/flax) or [Haiku](https://github.com/deepmind/haiku)? The main difference is that Equinox (a) offers a lot of advanced features not found in these libraries, like PyTree manipulation or runtime errors; (b) has a simpler way of building models: they're just PyTrees, so they can pass across JIT/grad/etc. boundaries smoothly._\n\n## Installation\n\nRequires Python 3.10+.\n\n```bash\npip install equinox\n```\n\nEquinox is also available through a community-supported build on [conda-forge](https://github.com/conda-forge/equinox-feedstock).\n\n## Documentation\n\nAvailable at [https://docs.kidger.site/equinox](https://docs.kidger.site/equinox).\n\n## Quick example\n\nModels are defined using PyTorch-like syntax:\n\n```python\nimport equinox as eqx\nimport jax\n\nclass Linear(eqx.Module):\n    weight: jax.Array\n    bias: jax.Array\n\n    def __init__(self, in_size, out_size, key):\n        wkey, bkey = jax.random.split(key)\n        self.weight = jax.random.normal(wkey, (out_size, in_size))\n        self.bias = jax.random.normal(bkey, (out_size,))\n\n    def __call__(self, x):\n        return self.weight @ x + self.bias\n```\n\nand are fully compatible with normal JAX operations:\n\n```python\n@jax.jit\n@jax.grad\ndef loss_fn(model, x, y):\n    pred_y = jax.vmap(model)(x)\n    return jax.numpy.mean((y - pred_y) ** 2)\n\nbatch_size, in_size, out_size = 32, 2, 3\nmodel = Linear(in_size, out_size, key=jax.random.PRNGKey(0))\nx = jax.numpy.zeros((batch_size, in_size))\ny = jax.numpy.zeros((batch_size, out_size))\ngrads = loss_fn(model, x, y)\n```\n\nFinally, there's no magic behind the scenes. All `eqx.Module` does is register your class as a PyTree. From that point onwards, JAX already knows how to work with PyTrees.\n\n## Citation\n\nIf you found this library to be useful in academic work, then please cite: ([arXiv link](https://arxiv.org/abs/2111.00254))\n\n```bibtex\n@article{kidger2021equinox,\n    author={Patrick Kidger and Cristian Garcia},\n    title={{E}quinox: neural networks in {JAX} via callable {P}y{T}rees and filtered transformations},\n    year={2021},\n    journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}\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[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[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  \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 Equinox](https://docs.kidger.site/equinox/awesome-list/)  \n[Awesome JAX](https://github.com/lockwo/awesome-jax): a longer list of other JAX projects.  \n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Fequinox","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpatrick-kidger%2Fequinox","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Fequinox/lists"}