[pallas:mosaic_gpu] Fix the tests following the changes to pl.core_map

PiperOrigin-RevId: 713283207

Author: superbobry

jax/_src/pallas/mosaic_gpu/core.py

@@ -483,7 +483,6 @@ class GPUMesh:
   # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.
   num_threads: int | None = None
   axis_names: tuple[str, ...] = ()
-  approx_math: bool = False
 
   def __post_init__(self):
     if len(self.axis_names) != len(self.grid) + (self.num_threads is not None):
@@ -521,21 +520,38 @@ def _gpu_mesh_discharge_rule(
     *args,
     mesh,
     jaxpr,
+    compiler_params,
+    interpret,
+    debug,
+    cost_estimate,
 ):
-  assert isinstance(mesh, GPUMesh)
+  if not isinstance(mesh, GPUMesh):
+    raise TypeError(f"Mesh must be a GPUMesh, got {type(mesh)}")
   if mesh.cluster:
     raise NotImplementedError
   if mesh.num_threads is None:
     raise NotImplementedError
+  if compiler_params and not isinstance(compiler_params, GPUCompilerParams):
+    raise TypeError(
+        "Compiler params must be a GPUCompilerParams, got"
+        f" {type(compiler_params)}"
+    )
+  if not compiler_params:
+    compiler_params = GPUCompilerParams()
   return pallas_core.default_mesh_discharge_rule(
       in_avals,
       out_avals,
       *args,
       jaxpr=jaxpr,
       grid=tuple(mesh.shape.items()),
       backend="mosaic_gpu",
-      compiler_params=GPUCompilerParams(approx_math=mesh.approx_math),
+      compiler_params=compiler_params,
+      debug=debug,
+      interpret=interpret,
+      cost_estimate=cost_estimate,
   )
+
+
 pallas_core._core_map_mesh_rules[GPUMesh] = _gpu_mesh_discharge_rule
 
 

jax/experimental/pallas/ops/gpu/attention_mgpu.py

@@ -200,9 +200,11 @@ def run(refs):
         grid=(batch_size, num_q_tiles, num_q_heads),
         num_threads=3,
         axis_names=("batch", "q_seq", "heads", "wg"),
-        approx_math=True,
     )
-    @pl.core_map(mesh)
+
+    @pl.core_map(
+        mesh, compiler_params=plgpu.GPUCompilerParams(approx_math=True)
+    )
     def _kernel_entry():
       compute_wgs = 2
       tiling = plgpu.TilingTransform((64, 64))

tests/pallas/mosaic_gpu_test.py

@@ -1481,14 +1481,11 @@ def copy_kernel(x_smem, o_smem, o_last_block_smem):
                                       index_map=lambda i, j: (0, 0))
                    ],
     )
-    mesh = plgpu.GPUMesh(
-        grid=(1,),
-        num_threads=3,
-        axis_names=("_", "wg",),
-        approx_math=True,
-    )
+    mesh = plgpu.GPUMesh(grid=(1,), num_threads=3, axis_names=("_", "wg"))
     def run(refs):
-      @pl.core_map(mesh)
+      @pl.core_map(
+          mesh, compiler_params=plgpu.GPUCompilerParams(approx_math=True)
+      )
       def _kernel_entry():
         pipeline(*refs)
     @jax.jit
@@ -1535,13 +1532,12 @@ def tiled_add_kernel(x_smem, y_smem, o_smem):
                 transforms=[])],
     )
     mesh = plgpu.GPUMesh(
-        grid=(1,),
-        num_threads=num_compute_wgs + 1,
-        axis_names=("_", "wg",),
-        approx_math=True,
+        grid=(1,), num_threads=num_compute_wgs + 1, axis_names=("_", "wg")
     )
     def run(refs):
-      @pl.core_map(mesh)
+      @pl.core_map(
+          mesh, compiler_params=plgpu.GPUCompilerParams(approx_math=True)
+      )
       def _kernel_entry():
         pipeline(*refs)
     @jax.jit