Fix cache init when JAX Array is created early (#25768)

jax/_src/compilation_cache.py

@@ -131,6 +131,11 @@ def _initialize_cache() -> None:
   with _cache_initialized_mutex:
     if _cache_initialized:
       return
+
+    path: str | None = config.compilation_cache_dir.value
+    # If the path is not set, the cache will not be built.
+    if not path:
+      return
     _cache_initialized = True
 
     # Nothing to do if the cache is disabled.
@@ -146,10 +151,6 @@ def _initialize_cache() -> None:
 
     global _cache
     assert _cache is None, "The cache has already been initialized!"
-    path: str | None = config.compilation_cache_dir.value
-    # If the path is not set, the cache will not be enabled.
-    if not path:
-      return
 
     cache_and_path = get_file_cache(path)
     if cache_and_path is None:

tests/compilation_cache_test.py

@@ -198,6 +198,30 @@ def test_jit(self):
     f(1.0)
     self.assertEqual(count_cache_items(), 2)
 
+
+  def test_change_compilation_cache_dir_reset_cache(self):
+    original_value = config.compilation_cache_dir.value
+    try:
+      cc.reset_cache()
+      self.assertFalse(cc._cache_initialized)
+      self.assertFalse(cc.is_persistent_cache_enabled())
+
+      a = jnp.zeros((2,3))
+      self.assertFalse(cc._cache_initialized)
+      self.assertFalse(cc.is_persistent_cache_enabled())
+
+      jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
+      @jax.jit
+      def f(x):
+        return x * 10
+      f(a)
+      self.assertTrue(cc._cache_initialized)
+      self.assertTrue(cc.is_persistent_cache_enabled())
+
+    finally:
+      config.update("jax_compilation_cache_dir", original_value)
+
+
   def test_xla_autofdo_profile_version(self):
     original_profile_version = config.jax_xla_profile_version.value
     with config.jax_xla_profile_version(original_profile_version + 1):