All Categories → No Category → jax

Top 42 jax open source projects

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

Deep learning operations reinvented (for pytorch, tensorflow, jax and others)

✭ 4,022

python deep-learning pytorch tensorflow keras numpy tensor chainer gluon cupy jax

🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

✭ 2,422

python cython javascript Sass Cuda deep-learning machine-learning pytorch tensorflow nlp natural-language-processing artificial-intelligence functional-programming ai mxnet spacy machine-learning-library type-checking jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

✭ 15,579

python C++Jupyter Notebook Starlark shell c jax

A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX

✭ 2,108

python Jupyter Notebook machine-learning pytorch tensorflow keras adversarial-examples adversarial-attacks jax

Flax is a neural network library for JAX that is designed for flexibility.

✭ 2,447

python Jupyter Notebook jax

Trax — Deep Learning with Clear Code and Speed

✭ 6,666

python Jupyter Notebook shell deep-learning machine-learning reinforcement-learning numpy deep-reinforcement-learning transformer jax

Python code for "Machine learning: a probabilistic perspective" (2nd edition)

✭ 4,197

python Jupyter Notebook machine-learning tensorflow pytorch colab pml flax jupyter-notebooks pymc3 pyro jax numpyro

TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...

✭ 3,094

python ruby smalltalk NewLisp javascript shell machine-learning tensorflow data dataset numpy datasets jax

Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).

✭ 61

python machine-learning computer-vision deep-learning neural-networks resnet flax jax

A small library for creating and manipulating custom JAX Pytree classes

✭ 29

python shell jax pytree

Official code for "Maximum Likelihood Training of Score-Based Diffusion Models", NeurIPS 2021 (spotlight)

✭ 49

python flax normalizing-flows jax diffusion-models score-matching neural-odes score-based-generative-modeling neurips-2021 score-based-generative-models

Implementation of https://srush.github.io/annotated-s4

✭ 133

python Makefile deep-learning jax

Documentation:

✭ 103

python deep-neural-networks lasagne theano deep-learning tensorflow numpy dataset jax

Massively parallel rigidbody physics simulation on accelerator hardware.

✭ 1,208

Jupyter Notebook python javascript reinforcement-learning robotics physics-simulation jax

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.

✭ 229

Jupyter Notebook machine-learning tutorial deep-learning jupyter numpy flax haiku lax jax xla optax learn-jax

Bayesian Optimization in JAX

✭ 24

python automatic-differentiation bayesian-optimization gaussian-process-regression jax

ML-Optimizers-JAX

Toy implementations of some popular ML optimizers using Python/JAX

✭ 37

python Jupyter Notebook machine-learning momentum gradient-descent optimization-algorithms adam adam-optimizer jax optimizers

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

✭ 50

python Makefile pde discretization jax differentiable-programming

Functional models and algorithms for sparse signal processing

✭ 38

Jupyter Notebook python functional-programming sparse-linear-systems convex-optimization linear-operators compressive-sensing wavelets sparse-representations jax l1-regularization first-order-conic-solvers

KoCLIP: Korean port of OpenAI CLIP, in Flax

✭ 80

Jupyter Notebook python flax jax roberta vision-transformer openai-clip

👩 Pytorch and Jax code for the Madam optimiser.

✭ 46

Jupyter Notebook python shell deep-learning optimizer pytorch neural-networks jax madam

Federated posterior averaging implemented in JAX

✭ 38

Jupyter Notebook python federated-learning jax

Contains code for the paper "Vision Transformers are Robust Learners" (AAAI 2022).

✭ 78

Jupyter Notebook python computer-vision tensorflow transformers pytorch robustness self-attention jax

chef-transformer

Chef Transformer 🍲 .

✭ 29

python Jupyter Notebook CSS recipes transformers flax recipe-generator jax huggingface t5 recipe-generation google-t5

A Python library for machine-learning and feedback loops on streaming data

✭ 36

python Makefile shell machine-learning reinforcement-learning time-series pandas xarray data-streaming jax

Factor graphs and nonlinear optimization for JAX

✭ 124

python factor-graphs nonlinear-optimization jax

PIX is an image processing library in JAX, for JAX.

✭ 271

python shell machine-learning image computer-vision image-processing jax

efficientnet-jax

EfficientNet, MobileNetV3, MobileNetV2, MixNet, etc in JAX w/ Flax Linen and Objax

✭ 114

python Dockerfile flax mixnet tpu jax mobilenetv2 mobilenetv3 efficientnet objax flax-linen

ShinRL: A Library for Evaluating RL Algorithms from Theoretical and Practical Perspectives (Deep RL Workshop 2021)

✭ 30

Jupyter Notebook python reinforcement-learning jax

A didactic Gaussian process package for researchers in Jax.

✭ 159

python TeX Makefile machine-learning bayesian-inference gaussian-processes jax

uvadlc notebooks

Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2022/Spring 2022

✭ 901

Jupyter Notebook tutorial deep-learning tutorials pytorch flax jax pytorch-lightning optax uvadlc

JAX-based Spiking Neural Network framework

✭ 60

python spiking-neural-networks jax deep-spiking-neural-networks biological-neural-networks spike-based-backpropagation

A GPT, made only of MLPs, in Jax

✭ 53

python deep-learning artificial-intelligence language-model multilayer-perceptron jax

revisiting rainbow

Revisiting Rainbow

✭ 71

Jupyter Notebook python ai tensorflow ml deep-reinforcement-learning rainbow rl jax

Unofficial JAX implementations of deep learning research papers

✭ 108

python machine-learning computer-vision deep-learning transformers artificial-intelligence convolutional-neural-networks flax jax research-paper-implementation

Computational Fluid Dynamics in JAX

✭ 399

Jupyter Notebook python cfd jax

JAX implementations of core Deep RL algorithms

✭ 61

python Makefile reinforcement-learning deep-learning deep-reinforcement-learning flax deepmind mpo sac actor-critic maximum-a-posteriori-estimation mujoco jax td3 soft-actor-critic

parallel-non-linear-gaussian-smoothers

Companion code in JAX for the paper Parallel Iterated Extended and Sigma-Point Kalman Smoothers.

✭ 17

Jupyter Notebook python parallel unscented-kalman-filter kalman-filter extended-kalman-filters state-space-model jax

ADAM implements a collection of algorithms for calculating rigid-body dynamics in Jax, CasADi, PyTorch, and Numpy.

✭ 51

python reinforcement-learning robotics optimization numpy automatic-differentiation motion-planning pytorch urdf rigid-body-dynamics jax casadi

JAX code for the paper "Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation"

✭ 43

Jupyter Notebook python machine-learning reinforcement-learning deep-learning automatic-differentiation deep-reinforcement-learning dqn gym flax haiku sac model-based-rl jax model-based-reinforcement-learning soft-actor-critic dm-haiku

Graphsignal Python agent

✭ 158

python shell debugging machine-learning monitoring deep-learning tensorflow inference pytorch artificial-intelligence openai jax inference-optimization huggingface onnxruntime langchain

1-42 of 42 jax projects