[Mosaic GPU][NFC] Split LaunchContext into a separate file

This will be used by multiple files in a follow up.

PiperOrigin-RevId: 707630113

Author: dimitar-asenov

jax/experimental/mosaic/gpu/__init__.py

@@ -19,16 +19,19 @@
 from .core import (
     Barrier as Barrier,
     ClusterBarrier as ClusterBarrier,
-    LaunchContext as LaunchContext,
-    MemRefTransform as MemRefTransform,
     TMABarrier as TMABarrier,
     ThreadSemantics as ThreadSemantics,
-    TileTransform as TileTransform,
-    TransposeTransform as TransposeTransform,
     Union as Union,
     as_gpu_kernel as as_gpu_kernel,
 )
 
+from .launch_context import (
+    LaunchContext as LaunchContext,
+    MemRefTransform as MemRefTransform,
+    TileTransform as TileTransform,
+    TransposeTransform as TransposeTransform,
+)
+
 if dialect is not None:
   from .dialect_lowering import (
       gpu_address_space_to_nvptx as gpu_address_space_to_nvptx,