Disable failing test cases when JAX_ENABLE_X64=1 in the Bazel CPU build

PiperOrigin-RevId: 705635799

Author: nitins17

tests/lax_autodiff_test.py

@@ -27,6 +27,7 @@
 import jax
 from jax import dtypes
 from jax import lax
+from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.util import NumpyComplexWarning
 from jax.test_util import check_grads
@@ -205,6 +206,8 @@ class LaxAutodiffTest(jtu.JaxTestCase):
   ))
   def testOpGrad(self, op, rng_factory, shapes, dtype, order, tol):
     rng = rng_factory(self.rng())
+    if jtu.test_device_matches(["cpu"]) and (op is lax.cosh or op is lax.cbrt) and config.enable_x64.value:
+      raise SkipTest("cosh and cbrt grad fails in x64 mode on CPU")  # b/383756018
     if jtu.test_device_matches(["tpu"]):
       if op is lax.pow:
         raise SkipTest("pow grad imprecise on tpu")

tests/lax_numpy_reducers_test.py

@@ -342,6 +342,8 @@ def np_fun(x):
   ))
   def testReducerPromoteInt(self, name, rng_factory, shape, dtype, axis,
                             keepdims, initial, inexact, promote_integers):
+    if jtu.test_device_matches(["cpu"]) and name == "sum" and config.enable_x64.value and dtype == np.float16:
+      raise unittest.SkipTest("sum op fails in x64 mode on CPU with dtype=float16")  # b/383756018
     np_op = getattr(np, name)
     jnp_op = getattr(jnp, name)
     rng = rng_factory(self.rng())