The Unified AI Framework
-
Updated
Jun 24, 2024 - Python
The Unified AI Framework
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
Clean single-file implementation of offline RL algorithms in JAX
Flax is a neural network library for JAX that is designed for flexibility.
A differentiable physics engine and multibody dynamics library for control and robot learning.
Higher-Order Interactions
PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper and Ada GPUs, to provide better performance with lower memory utilization in both training and inference.
NAACL '24 (Demo) / MlSys @ NeurIPS '23 - RedCoast: A Lightweight Tool to Automate Distributed Training and Inference
GPU-accelerated Evolutionary Multiobjective Optimization Using Tensorized RVEA.
Modular reinforcement learning library (on PyTorch and JAX) with support for NVIDIA Isaac Gym, Omniverse Isaac Gym and Isaac Lab
PEP 503 repository index for jax[cuda]
Really Fast End-to-End Jax RL Implementations
A library for scientific machine learning and physics-informed learning
Orbax provides common utility libraries for JAX users.
Add a description, image, and links to the jax topic page so that developers can more easily learn about it.
To associate your repository with the jax topic, visit your repo's landing page and select "manage topics."