JAX playground

Playground for learning jax in python

JAX

JAX is a library for high-performance numerical computing in python with autograd and XLA support.

Setup

1. Install miniconda
2. Create conda environment using conda env create -f environment.yml
3. Activate conda environment using conda activate jax-playground