#jax

Generative modeling for JAX/Flax NNX: VAEs, GANs, DDPM/DDIM/score/DiT/latent diffusion, normalizing flows, energy-based, autoregressive and geometric (SE(3), protein, point cloud, mesh) models; image, text, audio, molecular, protein, tabular and time-series modalities; a shared Trainer with callbacks and extensions; evaluation through calibrax.

Hardware accelerated, batchable and differentiable optimizers in JAX.

Sequential inference for state-space models in JAX: Kalman-family filters and smoothers, and SMC methods including particle filters, tempered SMC, and SMC².

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

Probabilistic reasoning and statistical analysis in TensorFlow

JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging.

Circulax is a differentiable circuit simulation framework built on JAX, Optimistix and Diffrax

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

JAX-based tensor network library with label-based contraction and symmetry-aware block-sparse tensors

Next-gen AI-native tensor-network-based quantum software framework

No repository description published.

No repository description published.

A JAX backend for Apple Metal Performance Shaders (MPS), enabling GPU-accelerated JAX computations on Apple Silicon.

e-Stellarator Simulation and Optimization Suite

No repository description published.

SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms

🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

Differentiable biophysical modeling in JAX (SAXS, NMR, CD), with Jupyter Notebook examples.

JAX version of STELLOPT's BOOZ_XFORM code