gordicaleksa / get-started-with-JAX
The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.
☆657Updated 11 months ago
Related projects ⓘ
Alternatives and complementary repositories for get-started-with-JAX
- JAX - A curated list of resources https://github.com/google/jax☆1,541Updated 3 months ago
- 100 exercises to learn JAX☆567Updated 2 years ago
- Optax is a gradient processing and optimization library for JAX.☆1,687Updated this week
- ☆785Updated last week
- ☆386Updated 3 weeks ago
- ☆535Updated last month
- Hardware accelerated, batchable and differentiable optimizers in JAX.☆930Updated last month
- All about the fundamental blocks of TF and JAX!☆271Updated 2 years ago
- CLU lets you write beautiful training loops in JAX.☆321Updated 2 months ago
- ML Collections is a library of Python Collections designed for ML use cases.☆893Updated 3 months ago
- functorch is JAX-like composable function transforms for PyTorch.☆1,394Updated this week
- ☆749Updated last month
- Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/☆2,103Updated last week
- PIX is an image processing library in JAX, for JAX.☆388Updated this week
- For optimization algorithm research and development.☆408Updated this week
- TensorDict is a pytorch dedicated tensor container.☆832Updated this week
- A Graph Neural Network Library in Jax☆1,374Updated 7 months ago
- Code for our NeurIPS 2022 paper☆361Updated last year
- ☆1,258Updated last month
- Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.☆1,402Updated 3 months ago
- Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/