Avoid assuming that jnp.sin will be traced in abstract mesh tests

The test does not clear the JAX caches, and jax.sin is a jitted closure
that's shared between all test methods, so there's no guarantee that someone
hasn't already traced sine at that same shape before. This only shows up rarely
since it depends on the subset of tests assigned to the same test executor.

PiperOrigin-RevId: 706706380

Author: apaszke

tests/pjit_test.py

@@ -4523,7 +4523,7 @@ def test_different_devices_wsc_abstract_mesh_cache_hit(self):
     def f(x):
       x = with_sharding_constraint(
           x, NamedSharding(mesh_lib.AbstractMesh(mesh1.shape_tuple), P('x')))
-      return jnp.sin(x)
+      return jax.lax.sin(x)
 
     with (
         jtu.count_jit_tracing_cache_miss() as tracing_count,
@@ -4536,7 +4536,8 @@ def f(x):
       # same num_devices but different devices.
       b = jax.device_put(out_a, NamedSharding(mesh2, P()))
       f(b)  # tracing and lowering cache *hit*
-    self.assertEqual(tracing_count(), 2)  # 1 miss for `f` and 1 miss for `sin`
+
+    self.assertEqual(tracing_count(), 1)
     self.assertEqual(lowering_count(), 1)
     self.assertEqual(compilation_count(), 2)  # 2 misses since devices differ.
 

tests/shard_map_test.py

@@ -811,7 +811,7 @@ def test_different_devices_shmap_abstract_mesh_cache_hit(self):
     def f(x):
       x = shard_map(lambda x: x, mesh=abstract_mesh, in_specs=P('i'),
                     out_specs=P('i'))(x)
-      return jnp.sin(x)
+      return jax.lax.sin(x)
 
     with (
         jtu.count_jit_tracing_cache_miss() as tracing_count,
@@ -825,7 +825,7 @@ def f(x):
       b = jax.device_put(out_a, NamedSharding(mesh2, P()))
       f(b)  # tracing and lowering cache *hit*
 
-    self.assertEqual(tracing_count(), 2)  # 1 miss for `f` and 1 miss for `sin`
+    self.assertEqual(tracing_count(), 1)
     self.assertEqual(lowering_count(), 1)
     self.assertEqual(compilation_count(), 2)  # 2 misses since devices differ.