41 Open source projects that are alternatives of or similar to ShinRL

A didactic Gaussian process package for researchers in Jax.

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

JAX-based Spiking Neural Network framework

A GPT, made only of MLPs, in Jax

Revisiting Rainbow

Unofficial JAX implementations of deep learning research papers

Computational Fluid Dynamics in JAX

JAX implementations of core Deep RL algorithms

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

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

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

Graphsignal Python agent

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

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

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

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

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

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

Trax — Deep Learning with Clear Code and Speed

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

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

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

A small library for creating and manipulating custom JAX Pytree classes

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

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

Documentation:

Massively parallel rigidbody physics simulation on accelerator hardware.

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.

Bayesian Optimization in JAX

Toy implementations of some popular ML optimizers using Python/JAX

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

Functional models and algorithms for sparse signal processing

KoCLIP: Korean port of OpenAI CLIP, in Flax

👩 Pytorch and Jax code for the Madam optimiser.

Federated posterior averaging implemented in JAX

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

Chef Transformer 🍲 .

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

Factor graphs and nonlinear optimization for JAX

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

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