[Pallas TPU] Add lowering for lax.cos_p

Fixes #24026 PiperOrigin-RevId: 680754948
jax-ml · Sep 30, 2024 · a24420e · a24420e
1 parent 23ce5a1
commit a24420e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
@@ -2001,6 +2001,13 @@ def _sin_lowering_rule(ctx: LoweringRuleContext, x):
 lowering_rules[lax.sin_p] = _sin_lowering_rule
 
 
+def _cos_lowering_rule(ctx: LoweringRuleContext, x):
+  return math.CosOp(x).result
+
+
+lowering_rules[lax.cos_p] = _cos_lowering_rule
+
+
 def _tanh_lowering_rule(ctx: LoweringRuleContext, x):
   return math.TanhOp(x).result
 

diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py
@@ -783,7 +783,7 @@ def test_elementwise(self, fn, dtype):
     # TODO(b/370578663): implement these lowerings on TPU
     if jtu.test_device_matches(["tpu"]) and fn in (
         jnp.acos, jnp.acosh, jnp.asin, jnp.asinh, jnp.atan, jnp.atanh,
-        jnp.cbrt, jnp.ceil, jnp.cos, jnp.cosh, lax.clz, jnp.expm1,
+        jnp.cbrt, jnp.ceil, jnp.cosh, lax.clz, jnp.expm1,
         jnp.floor, lax.population_count, jnp.sinh, jnp.tan,
     ):
       self.skipTest(f"{fn.__name__} not implemented on TPU")