[mgpu/pallas_mgpu] Pointwise tanh support

PiperOrigin-RevId: 700158250

Author: cperivol

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1179,6 +1179,12 @@ def _rsqrt_lowering_rule(ctx: LoweringRuleContext, x):
   [x_aval] = ctx.avals_in
   return _ensure_fa(x, x_aval.dtype).rsqrt(approx=ctx.module_ctx.approx_math)
 
+@register_lowering_rule(lax.tanh_p)
+def _tanh_lowering_rule(ctx: LoweringRuleContext, x):
+  [x_aval] = ctx.avals_in
+  return _ensure_fa(x, x_aval.dtype).tanh(approx=ctx.module_ctx.approx_math)
+
+
 @register_lowering_rule(lax.logistic_p)
 def _logistic_lowering_rule(ctx: LoweringRuleContext, x):
   [x_aval] = ctx.avals_in

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -918,6 +918,15 @@ def cos(self, *, approx: bool = False):
         self._lift_fast_unary("cos.approx.f32") if approx else mlir_math.cos
     )
 
+  def tanh(self, *, approx: bool = False):
+    if not ir.FloatType.isinstance(self.mlir_dtype):
+      raise NotImplementedError
+    if approx and self.mlir_dtype != ir.F32Type.get():
+      raise NotImplementedError
+    return self._pointwise(
+        self._lift_fast_unary("tanh.approx.f32") if approx else mlir_math.tanh
+    )
+
   def rsqrt(self, *, approx: bool = False):
     if not ir.FloatType.isinstance(self.mlir_dtype):
       raise NotImplementedError

tests/mosaic/gpu_test.py

@@ -1354,11 +1354,14 @@ def kernel(ctx, dst, _):
       ops=(
           (lambda x: -x, jax.lax.neg),
           (lambda x: x + 42, lambda x: x + 42),
+          (lambda x: x.tanh(), jax.lax.tanh),
       ),
       dtype=[jnp.float32, jnp.int32, jnp.uint32],
   )
   def test_unary(self, ops, dtype, m=64, n=32):
     op, np_op = ops
+    if np_op is jax.lax.tanh and jnp.issubdtype(dtype, jnp.integer):
+      raise self.skipTest("Tanh not supported for integer types")
 
     def kernel(ctx, dst, _):
       iota = iota_tensor(m, n, dtype)

tests/pallas/mosaic_gpu_test.py

@@ -70,8 +70,9 @@ class PallasCallTest(PallasTest):
       ("exp", jax.lax.exp),
       ("square", lambda x: x ** 2),
       ("rsqrt", jax.lax.rsqrt),
+      ("tanh", jax.lax.tanh, 1e-6),
   )
-  def test_unary_ops(self, unary):
+  def test_unary_ops(self, unary, rtol=1e-7):
     @functools.partial(
         pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
@@ -80,7 +81,7 @@ def kernel(x_ref, o_ref):
       o_ref[...] = unary(x_ref[...])
 
     x = jnp.arange(256).astype(jnp.float32)
-    np.testing.assert_array_equal(kernel(x), unary(x))
+    np.testing.assert_allclose(kernel(x), unary(x), rtol=rtol)
 
   def test_add_first(self):
     @functools.partial(