jax

03612554 - [Mosaic GPU] Add layout inference and initial lowering for `vector.{load,store}`.

[Mosaic GPU] Add layout inference and initial lowering for `vector.{load,store}`. For now, the lowering only works for the strided fragmented layout. This is mostly an exercise in plugging in lowering rules using `FragmentedArray`, and will be expanded shortly. PiperOrigin-RevId: 707031770