jax

71c76220 - [Pallas:MGPU] Change WG semantics convention to represent scalar arrays using scalars

[Pallas:MGPU] Change WG semantics convention to represent scalar arrays using scalars Previously every ShapedArray got converted to an MLIR vector which was more annoying than helpful. PiperOrigin-RevId: 730795455