jax

d36f6b3b - Update XLA dependency to use revision http://github.com/openxla/xla/commit/d776fe25fde3bb0f28485ccd7947ade00e3d3bf4