jax

525a13fb - [Mosaic GPU] Introduce `{Least,Most}ReplicatedExpression` in the equation system.

[Mosaic GPU] Introduce `{Least,Most}ReplicatedExpression` in the equation system. They will be used in layout inference, in order to propagate producer and consumer layouts. As a motivating example for `LeastReplicatedExpression`, think of an elementwise `add`: ``` v0 = producer0() # out_layouts = splat v1 = producer1() # out_layouts = strided r = add(v0, v1) ``` Looking at `add`, we will derive equations such that `v0 = v1 = r`. `v1` has a strided layout, which means that we can not make `add` use a `splat` layout, even though `v0` has a splat layout. Whenever we encounter such a case, we will want to use the least replicated layout. `MostReplicatedExpression` is necessary in order to handle consumers, in a similar fashion---with the thesis that following the flow of data, if the layout of a consumer has been determined, we know that the consumer can not be more replicated than its producer---and hence the most replicated consumer layout is also a safe choice. PiperOrigin-RevId: 778370106