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.

* Moves all the key classes down to C++ level, while keeping the APIs unchanged. Known small caveats: they are no longer Python dataclasses, and pattern matching might make pytype unhappy.

* Registered defaultdict and ordereddict via the keypath API now.

PiperOrigin-RevId: 694219933

chex/_src/dataclass.py

@@ -279,10 +279,10 @@ def _flatten_with_path(dcls):
   path = []
   keys = []
   for k, v in sorted(dcls.__dict__.items()):
+    keys.append(k)  # generate same aux data as flatten without path
     k = jax.tree_util.GetAttrKey(k)
     path.append((k, v))
-    keys.append(k)
-  return path, keys
+  return path, tuple(keys)
 
 
 @functools.cache