Distrax, but in equinox. Lightweight JAX library of probability distributions and bijectors.
-
Updated
Jun 12, 2024 - Python
Distrax, but in equinox. Lightweight JAX library of probability distributions and bijectors.
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
This is a JAX/Flax-based transformer language model trained on a Japanese dataset. It is based on the official Flax example code (lm1b).
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
🔥 A tool for visualizing and tracking your machine learning experiments. This repo contains the CLI and Python API.
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.
NGC-Learn: Neurobiological Learning and Biomimetic Systems Simulation in Python
A retargetable MLIR-based machine learning compiler and runtime toolkit.
FLS point cloud registration library.
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.
JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework
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.
Galactic and Gravitational Dynamics in Python (+ GPU and autodiff)
Code I used for my YouTube videos
Orbax provides common utility libraries for JAX users.
Flax is a neural network library for JAX that is designed for flexibility.
PEP 503 repository index for jax[cuda]
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."