jax

d08c35fb - [Mosaic GPU] Add an initial skeleton for an equation system for layout/transform inference.

[Mosaic GPU] Add an initial skeleton for an equation system for layout/transform inference. The idea will be to build an `EquationSystem` that describes how the layouts of the operands/results of MLIR operations relate to each other. Eventually, we'll be able to use an equation system to connect layouts in different memory spaces---e.g. `transforms` in SMEM, layouts in registers, and even TMEM layouts. PiperOrigin-RevId: 777562855