jax

9d421c91 - Plumb layout through the creation of PjRtArrays.

Plumb layout through the creation of PjRtArrays. This is in preparation to support arrays with no local shards, so that layout may not be accessible from a buffer. PiperOrigin-RevId: 730469597