jax

ab090e21 - Merge branch 'master' of https://github.com/google/jax into morejets