PyPI
avitai/artifex
A research-focused modular generative modeling library built on JAX/Flax NNX
MITAug 29, 2026 · metrics 2.10.0
All tags
Catalogue tag
#jax
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.
46 records
PyPI
tensorcircuit/tensorcircuit-ng
Next-gen AI-native tensor-network-based quantum software framework
Apache-2.0Aug 22, 2026 · metrics 2.10.0
PyPI
google-deepmind/distrax
No repository description published.
Apache-2.0Jul 30, 2026 · metrics 2.10.0
PyPI
araffin/sbx
SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms
MITSep 6, 2026 · metrics 2.10.0
PyPI · npm
explosion/thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
MITAug 27, 2026 · metrics 2.10.0
PyPI
tenax-lab/tenax
JAX-based tensor network library with label-based contraction and symmetry-aware block-sparse tensors
Apache-2.0Jul 19, 2026 · metrics 2.10.0
PyPI
google/jaxopt
Hardware accelerated, batchable and differentiable optimizers in JAX.
Apache-2.0Aug 4, 2026 · metrics 2.10.0
PyPI
OpenHCSDev/PolyStore
No repository description published.
MITAug 22, 2026 · metrics 2.10.0
PyPI
elkins-lab/diff-biophys
Differentiable biophysical modeling in JAX (SAXS, NMR, CD), with Jupyter Notebook examples.
MITJul 17, 2026 · metrics 2.10.0
PyPI
ai-hypercomputer/jetstream
JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).
Apache-2.0Jul 18, 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.
MITJul 17, 2026 · metrics 2.10.0
PyPI
HighDimensionalEconLab/nlls_gram
Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax NNX models
MITJul 18, 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.
MITJul 17, 2026 · metrics 2.10.0
PyPI
google/jetstream
JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).
Apache-2.0Jul 18, 2026 · metrics 2.10.0
PyPI
42Ar/jaxon
Save pytrees efficiently in hdf5 files.
GPL-3.0Aug 11, 2026 · metrics 2.10.0
PyPI
google-deepmind/jraph
A Graph Neural Network Library in Jax
Apache-2.0Aug 13, 2026 · metrics 2.10.0