58 dependents
| Package | Description | Downloads/month |
|---|---|---|
| Flax is a neural network library for JAX that is designed for flexibility. | 5.4M | |
| Orbax provides common checkpointing and persistence utilities for JAX users | 895K | |
| Package of Pathways-on-Cloud utilities | 410K | |
| Massively parallel rigidbody physics simulation on accelerator hardware. | 53K | |
| export JAX to ONNX | 37K | |
| Pax is a Jax-based machine learning framework for training large scale models. P... | 29K | |
| Modular, scalable library to train ML models | 24K | |
| Gemma open-weight LLM library, from Google DeepMind | 19K | |
| 17K | ||
| 12K | ||
| A Differentiable Data Pipeline Framework for JAX | 3K | |
| A framework for Alberta Plan research building the foundations of Continual AI. | 3K | |
| Use bacpipe to streamline the process of generating embeddings and analysing you... | 3K | |
| A library for fast Jax experimentation | 3K | |
| Foundation models in JAX/Flax | 3K | |
| A machine learning library intended for surrogate modeling tasks. | 3K | |
| Nasy's Deep Learning Toolkit | 2K | |
| A research-focused modular generative modeling library built on JAX/Flax NNX | 1K | |
| Continuous-discrete dynamical systems with JAX and related libraries. | 1K | |
| Finetuning utilities for AlphaGenome with custom heads and parameter freezing | 1K | |
| Hax - A suite for Heterogeneity analysis for CryoEM data | 1K | |
| JAX library for neural operator and PDE foundation model training. | 1K | |
| NumPyro integration for dynamical systems | 1K | |
| Training of CLIP in JAX | 1K | |
| A flexible and performant framework for training machine learning potentials. | 973 | |
| A simple JAX-powered simulation library for numerical experiments of fuzzy dark ... | 913 | |
| Opifex Computational Framework - JAX-native platform for scientific machine lear... | 903 | |
| Library for efficient training and application of Machine Learning Interatomic P... | 894 | |
| Reinforcement learning on general 2D physics environments in JAX. ICLR 2025 Oral... | 850 | |
| Natural language-driven JAX/Flax model building powered by Gemini | 711 | |
| A Simplistic trainer for Flax | 549 | |
| Wikidata and Wiktionary language data extraction | 546 | |
| Spark: Modular Spiking Neural Networks | 501 | |
| This is a Python library for solving Hamilton-Jacobi-Bellman (HJB) equations usi... | 463 | |
| A differentiable astrophysics simulator in JAX | 420 | |
| Library for machine learning on physical tensors | 419 | |
| Just Layers over JAX | 377 | |
| DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Ar... | 347 | |
| Statistical inference library built on JAX pytrees and optimistix optimizers | 333 | |
| Generative Models for Simulation-Based Inference in JAX | 326 | |
| Deep Learning for Bayesian Inference (DL4BI) | 321 | |
| LLM post-training library | 319 | |
| BEASTsim (BEnchmarking and Analysis of Spatial Transcriptomics simulations) is a... | 287 | |
| A package for GRAPE with feedback | 281 | |
| Learn spatially informed Waddington-like potentials for single-cell gene express... | 242 | |
| An advanced neural network framework with interpretability, generalization, robu... | 214 | |
| A lightspeed PIV images generator using JAX | 214 | |
| Library for machine learning on physical tensors | 148 | |
| Modelo de lenguaje avanzado basado en SSM y tecnologías experimentales | 147 | |
| Library for machine learning on physical tensors | 139 |
