jax

f2d16874 - [Pallas/Mosaic GPU] Carry `BlockSpec` transforms when lowering with warpgroup semantics.

[Pallas/Mosaic GPU] Carry `BlockSpec` transforms when lowering with warpgroup semantics. The `aval`s associated with refs that carry `BlockSpec` transforms are transformed prior to being fed into the `jaxpr`. Under warpgroup semantics, we want to generate IR that operates on untransformed references. This means that: 1. whenever we encounter a ref on which `BlockSpec` transforms have been applied, we need to reverse their effect before applying other transforms on the ref; 2. whenever we need to feed a ref into a `jaxpr` ourselves, we need to reapply these transforms in order to satisfy the typing expectations of the `jaxpr`: 3. no operation on a ref whose type has been transformed must remain visible in the generated IR. We achieve 1. and 2. by applying the transforms and their reverse transforms in the relevant places, and rely on `mosaic_gpu.reinterpret_cast` to convert inbetween types in the IR. We achieve 3. by relying on canonicalization rules centered on `mosaic_gpu.reinterpret_cast`, that allow all manifestations of the op to disappear. PiperOrigin-RevId: 901309524