Fix structure caching for enable_flax=False case by storing relative file names in structure rather than absolute paths.

PiperOrigin-RevId: 475845023

Author: cpgaffney1

orbax/checkpoint/pytree_checkpoint_handler.py

@@ -207,7 +207,7 @@ def _get_param_info(leaf):
       tspec = None
     elif isinstance(leaf, dict):
       tspec = None
-    elif isinstance(leaf, epath.Path):
+    elif isinstance(leaf, utils.Leaf):
       # Leaf is a param name.
       path = os.fspath(directory / leaf)
       tspec = serialization.get_tensorstore_spec(path)

orbax/checkpoint/utils.py

@@ -82,34 +82,34 @@ def rmtree(path: epath.Path):
   path.rmdir()
 
 
-Leaf = epath.Path
+Leaf = str
 
 
 def pytree_structure(directory: Path) -> PyTree:
   """Reconstruct state dict from saved model format in `directory`."""
   directory = epath.Path(directory)
 
-  def add_nested_key(subtree, nested_key, full_name):
+  def add_nested_key(subtree, nested_key, key_name):
     if not nested_key:
       return subtree
 
     current = nested_key[0]
 
     if len(nested_key) == 1:
       assert current not in subtree
-      subtree[current] = full_name
+      subtree[current] = key_name
       return subtree
 
     subkeys = nested_key[1:]
     if current not in subtree:
       subtree[current] = {}
-    subtree[current] = add_nested_key(subtree[current], subkeys, full_name)
+    subtree[current] = add_nested_key(subtree[current], subkeys, key_name)
     return subtree
 
   keys = directory.iterdir()
   tree = {}
   for k in keys:
-    tree = add_nested_key(tree, k.name.split('.'), k)
+    tree = add_nested_key(tree, k.name.split('.'), k.name)
   return tree
 
 

setup.py

@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='orbax',
-    version='0.0.10',
+    version='0.0.11',
     description='Orbax',
     long_description=_LONG_DESCRIPTION,
     long_description_content_type='text/markdown',