jax

3e29c9fb - add initial jet tests

add initial jet tests Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> Co-authored-by: Dougal Maclaurin <dougalm@google.com>