Equinox Python Packages

equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

1.7M 3K 194

lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

311K 545 38

optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/

217K 580 48

diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

189K 2K 180

exponax

Efficient Differentiable n-d PDE Solvers in JAX.

16K 200 17

pdequinox

Neural Emulator Architectures in JAX.

16K 24 0

trainax

Neural Emulator Architectures in JAX.

15K 24 0

quax

Multiple dispatch over abstract array types in JAX.

8K 141 7

eqxvision

Root package info.

3K 111 15

esm2quinox

An implementation of ESM2 in Equinox+JAX

2K 36 4

diffraxtra

diffrax extras: OOP and vectorization

697 6 3

sbgm

Score-based Diffusion models in JAX.

627 18 1

oncequinox

micro-library for making Equinox singletons

398 1 2

lerax

Reinforcement Learning in Jax

243 6 0

mccube

Markov Chain Cubature via JAX. Efficient SDE solving and Bayesian inference.

238 4 0

parametrix

🪄 Flax-like computed parameters for bare JAX (and Equinox)

163 0 0

discretax

Discretax is a light weight collection of state space models implemented in JAX ⚡️

160 38 5

somap

A flexible, fast and scalable python library for Self-Organizing Maps

148 16 1

linax

Discretax is a light weight collection of state space models implemented in JAX ⚡️

92 38 5

feedbax

Optimal feedback control + interventions, in JAX.

79 6 1

katmer

katmer is a powerful library for optimizing the design of optical thin films using automatic differentiation via JAX and Equinox, enabling efficient and accurate inverse design solutions.

78 11 3

frontx

⚛️ Nonlinear diffusion problems with JAX

76 2 0