Fixed a typo in Pallas GPU lowering

`abs` is not available in `triton.compat.math`. PiperOrigin-RevId: 601709135
jax-ml · Jan 26, 2024 · f34bcc3 · f34bcc3
1 parent 8901552
commit f34bcc3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/pallas/triton/lowering.py b/jax/_src/pallas/triton/lowering.py
@@ -463,7 +463,7 @@ def _cumsum_lowering_rule(
 _TRITON_FN_MAPPING = {
     # Unary ops.
     lax.neg_p: tc.semantic.minus,
-    lax.abs_p: tc.math.abs,
+    lax.abs_p: tc.abs,
     lax.ceil_p: tc.math.ceil,
     lax.floor_p: tc.math.floor,
     lax.exp_p: tc.math.exp,