Export KeyPath and related types to jax.tree_util

These types lie on the APIs in `jax.tree_util`, so it makes sense to export them. PiperOrigin-RevId: 657987755
jax-ml · Jul 31, 2024 · a207fe9 · a207fe9
1 parent 9dba6eb
commit a207fe9
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/docs/jax.tree_util.rst b/docs/jax.tree_util.rst
@@ -26,6 +26,8 @@ List of Functions
    treedef_children
    treedef_is_leaf
    treedef_tuple
+   KeyEntry
+   KeyPath
    keystr
 
 Legacy APIs

diff --git a/jax/tree_util.py b/jax/tree_util.py
@@ -42,6 +42,8 @@
     DictKey as DictKey,
     FlattenedIndexKey as FlattenedIndexKey,
     GetAttrKey as GetAttrKey,
+    KeyEntry as KeyEntry,
+    KeyPath as KeyPath,
     Partial as Partial,
     PyTreeDef as PyTreeDef,
     SequenceKey as SequenceKey,

diff --git a/tests/package_structure_test.py b/tests/package_structure_test.py
@@ -28,11 +28,15 @@ def _mod(module_name: str, *, include: Sequence[str] = (), exclude: Sequence[str
 
 
 class PackageStructureTest(jtu.JaxTestCase):
+
   @parameterized.parameters([
-    # TODO(jakevdp): expand test to other public modules.
-    _mod("jax.errors"),
-    _mod("jax.nn.initializers"),
-    _mod("jax.tree_util", exclude=['PyTreeDef', 'default_registry']),
+      # TODO(jakevdp): expand test to other public modules.
+      _mod("jax.errors"),
+      _mod("jax.nn.initializers"),
+      _mod(
+          "jax.tree_util",
+          exclude=["PyTreeDef", "default_registry", "KeyEntry", "KeyPath"],
+      ),
   ])
   def test_exported_names_match_module(self, module_name, include, exclude):
     """Test that all public exports have __module__ set correctly."""