[Pallas TPU] Add lowering for lax.tan_p

This is a follow-up of #24028, which adds lowering for `lax.cos_p`

PiperOrigin-RevId: 681134519

jax/_src/pallas/mosaic/lowering.py

@@ -2020,6 +2020,13 @@ def _cos_lowering_rule(ctx: LoweringRuleContext, x):
 lowering_rules[lax.cos_p] = _cos_lowering_rule
 
 
+def _tan_lowering_rule(ctx: LoweringRuleContext, x):
+  return math.TanOp(x).result
+
+
+lowering_rules[lax.tan_p] = _tan_lowering_rule
+
+
 def _tanh_lowering_rule(ctx: LoweringRuleContext, x):
   return math.TanhOp(x).result
 

tests/pallas/ops_test.py

@@ -776,7 +776,7 @@ def test_elementwise(self, fn, dtype):
     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.cosh, lax.clz, jnp.expm1,
-        jnp.floor, lax.population_count, jnp.sinh, jnp.tan,
+        jnp.floor, lax.population_count, jnp.sinh,
     ):
       self.skipTest(f"{fn.__name__} not implemented on TPU")