jax

832eb2d8 - [Mosaic GPU] Allow unions nested inside the smem ref tree

[Mosaic GPU] Allow unions nested inside the smem ref tree We don't use this capability just yet, but I want to start allocating barriers as part of the scratch and this will push the unions deeper into the tree. PiperOrigin-RevId: 655475839