importing etil instead of tf

jax/experimental/compilation_cache/gfile_cache.py

@@ -15,31 +15,34 @@
 import pathlib
 
 from jax.experimental.compilation_cache.cache_interface import CacheInterface
-import tensorflow.io.gfile as gfile
+from etils import epath
 from absl import logging
 
 class GFileCache(CacheInterface):
 
   def __init__(self, path: str):
     """Sets up a cache at 'path'. Cached values may already be present."""
-    self._path = pathlib.Path(path)
-    gfile.makedirs(self._path)
+    self._path = epath.Path(path)
+    self._path.mkdir(parents=True, exist_ok=True)
 
   def get(self, key: str):
     """Returns None if 'key' isn't present."""
     if not key:
       raise ValueError("key cannot be empty")
-    path_to_key = str(self._path / key)
-    if gfile.exists(path_to_key):
-      with gfile.GFile(path_to_key, "rb") as file:
-        return file.read()
+    path_to_key = self._path / key
+    if path_to_key.exists():
+      return path_to_key.read_bytes()
     else:
       return None
 
   def put(self, key: str, value: bytes):
     """Adds new cache entry."""
     if not key:
       raise ValueError("key cannot be empty")
-    path_to_new_file = str(self._path / key)
-    with gfile.GFile(path_to_new_file, "w") as f:
-      f.write(value)
+    path_to_new_file = self._path / key
+    if 'gs://' in str(path_to_new_file):
+      path_to_new_file.write_bytes(value)
+    else:
+      tmp_path = epath.Path(str(self._path / "_") + key)
+      tmp_path.write_bytes(value)
+      tmp_path.rename(str(path_to_new_file))