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.
☆689Updated last year
Alternatives and similar repositories for get-started-with-JAX:
Users that are interested in get-started-with-JAX are comparing it to the libraries listed below
- JAX - A curated list of resources https://github.com/google/jax☆1,699Updated this week
- Optax is a gradient processing and optimization library for JAX.☆1,802Updated this week
- 100 exercises to learn JAX☆576Updated 2 years ago
- ☆823Updated last week
- ☆1,288Updated 3 weeks ago
- ML Collections is a library of Python Collections designed for ML use cases.☆917Updated last week
- ☆764Updated last month
- All about the fundamental blocks of TF and JAX!☆275Updated 3 years ago
- Hardware accelerated, batchable and differentiable optimizers in JAX.☆947Updated 5 months ago
- A Graph Neural Network Library in Jax☆1,402Updated 11 months ago
- Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/☆2,239Updated this week
- JAX-based neural network library☆2,960Updated 3 weeks ago
- TensorDict is a pytorch dedicated tensor container.☆879Updated this week
- PIX is an image processing library in JAX, for JAX.☆402Updated this week
- functorch is JAX-like composable function transforms for PyTorch.☆1,410Updated this week
- CLU lets you write beautiful training loops in JAX.☆333Updated 2 weeks ago
- Advanced evolutionary computation library built directly on top of PyTorch, created at NNAISENSE.☆1,038Updated this week
- Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/☆1,541Updated last week
- Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/☆1,322Updated this week
- ☆548Updated 2 months ago
- A PyTorch library entirely dedicated to neural differential equations, implicit models and related numerical methods☆1,442Updated 9 months ago
- A library to inspect and extract intermediate layers of PyTorch models.☆471Updated 2 years ago
- Code for our NeurIPS 2022 paper☆366Updated 2 years ago
- All about the fundamentals and working of Diffusion Models☆153Updated 2 years ago
- Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.☆1,417Updated 6 months ago
- Library for reading and processing ML training data.☆386Updated this week
- MLCommons Algorithmic Efficiency is a benchmark and competition measuring neural network training speedups due to algorithmic improvement…☆364Updated last week
- Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.☆567Updated this week
- For optimization algorithm research and development.☆491Updated this week
- Tensors, for human consumption☆1,183Updated 3 months ago