Merge pull request #22279 from ayaka14732:ayx/lowering/shift_right_ar…

…ithmetic PiperOrigin-RevId: 650272128
google · Jul 8, 2024 · 9405d46 · 9405d46
2 parents a88ee4e + f013bd5
commit 9405d46
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
@@ -2258,6 +2258,15 @@ def _shift_left_lowering_rule(ctx: LoweringRuleContext, x, d):
 skip_mlir_conversions.add(lax.shift_left_p)
 
 
+def _shift_right_arithmetic_lowering_rule(ctx: LoweringRuleContext, x, d):
+  x, d = _bcast(x, d, *ctx.avals_in, *ctx.avals_out)
+  return arith.ShRSIOp(x, d).result
+
+
+lowering_rules[lax.shift_right_arithmetic_p] = _shift_right_arithmetic_lowering_rule
+skip_mlir_conversions.add(lax.shift_right_arithmetic_p)
+
+
 def _shift_right_logical_lowering_rules(ctx: LoweringRuleContext, x, d):
   x, d = _bcast(x, d, *ctx.avals_in, *ctx.avals_out)
   return arith.ShRUIOp(x, d).result

diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py
@@ -88,6 +88,7 @@ def setUp(self):
           (lax.bitwise_or, jnp.int32),
           (lax.bitwise_xor, jnp.int32),
           (lax.shift_left, jnp.int32),
+          (lax.shift_right_arithmetic, jnp.int32),
           (lax.shift_right_logical, jnp.int32),
       ]
   )