jax

5a7325eb - [Mosaic GPU] Only allow extracting a non-replicated layout assignment out of a `LeastReplicatedExpression` very conservatively.

[Mosaic GPU] Only allow extracting a non-replicated layout assignment out of a `LeastReplicatedExpression` very conservatively. It's not correct to propagate replicated layouts upwards in the general case. PiperOrigin-RevId: 781431727