jax

84066b77 - [Mosaic GPU] Change layout inference tests to rely on explicit `layout_cast`s.

[Mosaic GPU] Change layout inference tests to rely on explicit `layout_cast`s. This uncovers a propagation bug, whereby opportunities to propagate replicated layouts that were not already explicitly annotated as attributes downwards would be missed, because layout propagation started with a backwards pass. We now changed the implementation to start with a forward pass. Some additional edits: 1. I changed the layout in `test_optimization_barrier_op_propagates_user_layouts`. Generally, propagating replicated layouts upwards is not a safe thing to do, and we should have properly caught that. The upcoming infrastructure will recognize such issues, so we don't bother attempting to fix the underlying problem here; 2. I got rid of `test_infer_layout_propagates_func_layouts_to_ops` since we no longer care about `FuncOp`s. This simplification will allow the new infrastructure to not concern itself with `FuncOp`s, on which we were putting inconsistent expectations, and which would add quite a bit of complexity. PiperOrigin-RevId: 773277630