jax

17a4531e - Update XLA dependency to use revision http://github.com/openxla/xla/commit/4ef447c041b91498bf04b30e9395f703facff9cf