jax

a2973be0 - Don't add `mhlo.layout_mode = "default"` since that is the default even in PJRT and will help reduce cruft in the IR