[Mosaic GPU] Add support for transforms in gather TMA
This is ultimately what we need, as we need the data to be tiled in order to feed it to MMA.
Still no support for indexing of individual dimensions, though it should not be hard at this point.
PiperOrigin-RevId: 796364358
