jax

e17672bd - explorations with @duvenaud and @dougalm