[JAX] Ignore xla_gpu_experimental_autotune_cache_mode when calculating module hash.

PiperOrigin-RevId: 698789020

Author: Google-ML-Automation

jax/_src/cache_key.py

@@ -21,6 +21,7 @@
 from typing import cast as type_cast
 
 from jax._src import config
+from jax._src.lib import version as jaxlib_version
 from jax._src.lib import version_str as jaxlib_version_str
 from jax._src.lib import xla_client
 from jax._src.lib.mlir import ir
@@ -225,6 +226,8 @@ def _hash_serialized_compile_options(hash_obj, compile_options_obj,
   debug_options.xla_dump_hlo_as_long_text = False
   debug_options.xla_dump_disable_metadata = False
   debug_options.xla_dump_hlo_pipeline_re = ""
+  if jaxlib_version > (0, 4, 35):
+    debug_options.xla_gpu_experimental_autotune_cache_mode = 0
   # Optional way to specify the cuda install path to be used by the compiler.
   # This could possibly affect the cuda version compiled with, but this should
   # already be included in the platform information (and might not be reflected

tests/cache_key_test.py

@@ -31,6 +31,7 @@
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.lib import xla_client
+from jax._src.lib import version as jaxlib_version
 from jax._src.lib.mlir import ir
 from jax._src.mesh import Mesh
 from jax._src.partition_spec import PartitionSpec as P
@@ -68,6 +69,8 @@ def test_serialized_compile_options(self):
     debug_options.xla_dump_hlo_as_long_text = True
     debug_options.xla_dump_disable_metadata = True
     debug_options.xla_dump_hlo_pipeline_re = "xyzzy"
+    if jaxlib_version > (0, 4, 35):
+      debug_options.xla_gpu_experimental_autotune_cache_mode = 2
     hash2 = self.get_hashed_value(
         cache_key._hash_serialized_compile_options, compile_options
     )