339 dependents
| Package | Description | Downloads/month |
|---|---|---|
| 4.4M | ||
| CLU lets you write beautiful training loops in JAX. | 830K | |
| Accurate Quantized Training library. | 581K | |
| a Jax quantization library | 494K | |
| 228K | ||
| Differentiable, Hardware Accelerated, Molecular Dynamics | 95K | |
| A Lightweight LLM Post-Training Library | 68K | |
| Massively parallel rigidbody physics simulation on accelerator hardware. | 53K | |
| RL Environments in JAX 🌍 | 48K | |
| export JAX to ONNX | 37K | |
| A JAX-native High Performance Eval Metrics Library | 35K | |
| Functionalities such as a layers for building neural networks in Jax. | 31K | |
| Pax is a Jax-based machine learning framework for training large scale models. P... | 29K | |
| Dopamine is a research framework for fast prototyping of reinforcement learning ... | 26K | |
| Modular, scalable library to train ML models | 24K | |
| Gemma open-weight LLM library, from Google DeepMind | 19K | |
| Evolution Strategies in JAX 🦎 | 19K | |
| 17K | ||
| JAX-accelerated meta-reinforcement learning environments inspired by XLand and M... | 16K | |
| Multi-Party Programming Language | 13K | |
| Machine learning algorithms for many-body quantum systems | 13K | |
| 12K | ||
| Transport surrogate models for Tokamak fusion. | 11K | |
| Simulation Based Inference of Cosmological parameters in Jax using type Ia super... | 10K | |
| Accelerate, Optimize performance with streamlined training and serving options w... | 9K | |
| Single-cell perturbation analysis | 7K | |
| Transformer acceleration library - Jax Lib | 7K | |
| ⚡ Flashbax: Accelerated Replay Buffers in JAX | 7K | |
| (Crafter + NetHack) in JAX. ICML 2024 Spotlight. | 7K | |
| Multi-Agent Reinforcement Learning with JAX | 6K | |
| A decoder-only Transformer built from scratch in JAX and Flax NNX | 5K | |
| JetStream is a throughput and memory optimized engine for LLM inference on XLA d... | 4K | |
| Accelerated Quality-Diversity | 4K | |
| Accelerated minigrid environments with JAX | 4K | |
| (EasyDel Former) is a utility library designed to simplify and enhance the devel... | 4K | |
| A JAX implementation of the Metalog distribution for flexible probability modeli... | 4K | |
| Cost constrained environments with a gymnax interface | 4K | |
| EvoJAX: Hardware-accelerated Neuroevolution. | 3K | |
| A Differentiable Data Pipeline Framework for JAX | 3K | |
| A framework for Alberta Plan research building the foundations of Continual AI. | 3K | |
| Behavioral sequencing and phenotyping with lightweight task specific adaptation | 3K | |
| Python library for modeling DFAs, Moore Machines, and Transition Systems in JAX. | 3K | |
| JAX SparseCore API | 3K | |
| python computations in science and engineering | 3K | |
| Use bacpipe to streamline the process of generating embeddings and analysing you... | 3K | |
| SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms | 3K | |
| Foundation models in JAX/Flax | 3K | |
| A versatile and easy to understand Diffusion library | 3K | |
| A machine learning library intended for surrogate modeling tasks. | 3K | |
| Unified benchmarking and profiling framework for the JAX scientific ML ecosystem... | 2K |
