Update lowering.py

jax-ml · Dec 17, 2024 · 472cda8 · 472cda8
1 parent 45598f2
commit 472cda8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/pallas/triton/lowering.py b/jax/_src/pallas/triton/lowering.py
@@ -1620,9 +1620,9 @@ def _broadcast_in_dim_lowering_rule(
 
 
 @register_lowering(lax.squeeze_p)
-def _squeeze_lowering_rule(ctx: LoweringRuleContext, a, *, dimensions, sharding=None):
+def _squeeze_lowering_rule(ctx: LoweringRuleContext, a, *, dimensions):
   del dimensions
-  return _reshape_lowering_rule(ctx, a, new_sizes=None, dimensions=None, sharding=sharding)
+  return _reshape_lowering_rule(ctx, a, new_sizes=None, dimensions=None, sharding=None)
 
 
 @register_lowering(lax.reshape_p)