Minor improvement to dataclass flatten_with_path.

PiperOrigin-RevId: 557824539
google-deepmind · Aug 17, 2023 · d5bd8e9 · d5bd8e9
1 parent 3addcce
commit d5bd8e9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/chex/_src/dataclass.py b/chex/_src/dataclass.py
@@ -264,6 +264,7 @@ def _flatten_with_path(dcls):
   path = []
   keys = []
   for k, v in sorted(dcls.__dict__.items()):
+    k = jax.tree_util.GetAttrKey(k)
     path.append((k, v))
     keys.append(k)
   return path, keys