Fix copy-paste typo

Google-ML-Automation · Google-ML-Automation · commit eb2dd2ab0f2e · 2024-11-07T15:02:16.000-08:00
PiperOrigin-RevId: 694259486
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
@@ -1914,7 +1914,7 @@ def _sub_lowering_rule(ctx: LoweringRuleContext, x, y):
 
 
 lowering_rules[lax.sub_p] = _sub_lowering_rule
-skip_mlir_conversions.add(lax.max_p)
+skip_mlir_conversions.add(lax.sub_p)
 
 
 def _mul_lowering_rule(ctx: LoweringRuleContext, x, y):
