C++ tree with path API

Make tree_util.tree_flatten_with_path and tree_map_with_path APIs to be C++-based, to speed up the pytree flattening.

PiperOrigin-RevId: 694219933

chex/_src/dataclass.py

@@ -282,7 +282,7 @@ def _flatten_with_path(dcls):
     k = jax.tree_util.GetAttrKey(k)
     path.append((k, v))
     keys.append(k)
-  return path, keys
+  return path, tuple([k.name for k in keys])
 
 
 @functools.cache