PyPI
google/orbax
Orbax provides common checkpointing and persistence utilities for JAX users
Apache-2.0Oct 21, 2026 · metrics 2.10.0
All tags
Catalogue tag
#flax
Every repository in the public record carrying this tag — from its GitHub topics or the keywords its package registries publish. Health is measured under the same versioned methodology as the rest of the record.
6 records
PyPI
avitai/artifex
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.
MITOct 21, 2026 · metrics 2.10.0
PyPI
avitai/substrax
JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging.
MITSep 26, 2026 · metrics 2.10.0
NuGet · npm
outfox/fennecs
...the little C# ECS that loves you back! 🦊
MITSep 25, 2026 · metrics 2.10.0
PyPI
elkins-lab/resonance-flow
JAX-native differentiable protein folding framework integrating experimental NMR constraints (RDCs) and biophysical "self-correction." Several Jupyter Notebooks visualize the concepts.
MITSep 9, 2026 · metrics 2.10.0
PyPI
elkins-lab/diff-ensemble
Differentiable VAE framework for predicting protein structural ensembles (IDPs) consistent with SAXS and NMR data. Built on JAX/Flax.
MITSep 9, 2026 · metrics 2.10.0