A unified scientific machine learning framework built on JAX/Flax NNX
-
Updated
May 21, 2026 - Python
#
flax-nnx
Here are 14 public repositories matching this topic...
Titans: Learning to Memorize at Test Time
-
Updated
Jan 23, 2026 - Python
nanoGPT for Diffusion Language Models
-
Updated
May 14, 2026 - Python
Unified benchmarking and profiling framework for the JAX scientific ML ecosystem. Timing, GPU/energy monitoring, FLOPS counting, roofline analysis, statistical testing, regression detection, and CI integration.
python benchmarking machine-learning deep-learning numpy ci-cd statistical-analysis scientific-computing flax profiling performance-testing jax energy-monitoring mlflow regression-detection flops wandb gpu-profiling flax-nnx roofline-analysis
-
Updated
May 4, 2026 - Python
Flax NNX implementation of common metrics.
-
Updated
Jan 17, 2026 - Python
End-to-end differentiable bioinformatics pipelines built on JAX/Flax NNX
machine-learning bioinformatics deep-learning genomics protein-structure molecular-dynamics computational-biology rna-structure drug-discovery variant-calling sequence-alignment single-cell-genomics differentiable jax multiomics foundation-models flax-nnx differentiable-bioinformatics-pipelines
-
Updated
May 2, 2026 - Python
A Differentiable Data Pipeline Framework for JAX
-
Updated
May 1, 2026 - Python
A research-focused modular generative modeling library built on JAX/Flax NNX
research generative-adversarial-network vae gans multi-modal normalizing-flows variational-autoencoder generative-models diffusion-model geometric-models geometric-deep-learning jax diffusion-models energy-based-models autoregressive-model generative-ai flax-nnx
-
Updated
May 5, 2026 - Python
Molecular active learning with JAX
-
Updated
Apr 30, 2026 - Python
KGE-JAXed: A simple knowledge graph embedding library created in JAX
-
Updated
May 25, 2026 - Python
JAX/Flax NNX port of Karpathy's nanochat, optimized for multi-host Cloud TPU pods.
-
Updated
May 15, 2026 - Python
A 2-in-1 notebook-based tutorial and implementation of Manifold-Constrained Hyper-Connections using JAX
-
Updated
Feb 23, 2026 - Jupyter Notebook
JaxNN: Foundation Models in JAX/Flax
-
Updated
May 12, 2026 - Python
Recax is an experimental JAX/Flax NNX recommender systems framework
-
Updated
May 27, 2026 - Python
Improve this page
Add a description, image, and links to the flax-nnx topic page so that developers can more easily learn about it.
Add this topic to your repo
To associate your repository with the flax-nnx topic, visit your repo's landing page and select "manage topics."