[Pallas:MGPU] Fix an overly strict precision requirement in tests

They started failing after we allowed LLVM to perform contractions of
adds and muls, but the difference is tiny.

PiperOrigin-RevId: 701961845

tests/pallas/mosaic_gpu_test.py

@@ -492,9 +492,7 @@ def layer_norm_np(x):
         jax.random.uniform(jax.random.key(42), shape=(256,), dtype=jnp.float32)
         * input_factor
     )
-    # TODO(cperivol): find out why in this particular case we have a small-ish error.
-    rtol = 1e-07 if input_factor > 10 else 5e-5
-    np.testing.assert_allclose(layer_norm(x), layer_norm_np(x), rtol=rtol)
+    np.testing.assert_allclose(layer_norm(x), layer_norm_np(x), rtol=5e-5)
 
   def test_print(self):
     @functools.partial(