jax

2da8d573 - [Mosaic GPU] Propagate layouts from consumers and producers in the new layout inference.

[Mosaic GPU] Propagate layouts from consumers and producers in the new layout inference. While input and output layouts on a single op are usually connected by absolute constraints (represented by `Equation`s), the same does not hold for layouts of values across different operations. Take the following program: ``` v = producer() # out_layouts = ? consumer(v) # in_layouts = ? ``` Here, `producer` has a single `out_layout`, which correspond to the layout of its result `v`. The subsequent `consumer` call has a single `in_layout`, which may correspond go `producer`'s `out_layout`, but may also be different---e.g. if both `producer` and `consumer` have different strict constraints on what their respective layout should be. The crux of the issue is that, while we most often want to propagate layouts without changing them as much as possible, it is always valid to relayout a `vector` inbetween ops, making any constraint involving several ops a "soft" constraint. We introduce `Hint`s as a way to represent these soft constraints. `Hint`s are allow introducing new, likely-to-be-good assumptions in underspecified equation systems. Also add a layout inference rule for `layout_cast` in order to enable a new test. PiperOrigin-RevId: 778876118