Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
-
Updated
Jun 10, 2024 - Python
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
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.
Galactic and Gravitational Dynamics in Python (+ GPU and autodiff)
Coordinates in JAX
Loopy belief propagation for factor graphs on discrete variables in JAX
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
Collection of eclectic utils for python.
Algorithms for inverse design
This repository has code for the paper Bayesian prior impact assessment for dynamical systems described by ordinary differential equations
A retargetable MLIR-based machine learning compiler and runtime toolkit.
Orbax provides common utility libraries for JAX users.
Distrax, but in equinox. Lightweight JAX library of probability distributions and bijectors.
Custom types for topology optimization
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
Deep Learning for humans
Automated Machine Learning on Kubernetes
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."