[pallas:mosaic_gpu] Do not DCE the jaxpr in the lowering pass

There isn't an obvious reason for doing DCE there. PiperOrigin-RevId: 680534567
jax-ml · Sep 30, 2024 · b3fca90 · b3fca90
1 parent 21fea5b
commit b3fca90
Showing 1 changed file with 0 additions and 5 deletions.
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -239,11 +239,6 @@ def lower_jaxpr_to_module(
         "Only Blocked indexing mode is supported in Mosaic GPU lowering."
     )
 
-  with grid_mapping.trace_env():
-    jaxpr, _ = pe.dce_jaxpr(
-        jaxpr, [True] * len(jaxpr.outvars), instantiate=True
-    )
-
   block = (128, 1, 1)
   params = compiler_params.get("mosaic_gpu", {})
   approx_math = params.get("approx_math", False)