[Mosaic GPU] Use explicit load/store methods instead of __getitem__/__setitem__
We pretty never use slicing in those methods and I want to add the ability to load
in other layouts than the default one (which means we will need extra non-index arguments).
PiperOrigin-RevId: 756707034