Repository: jax-ml/jax
Pull request #28300 (2025-04-25 21:57:01)

Add is_leaf_with_path predicate.
Add is_leaf_with_path predicate.

This is an alternative of is_leaf that allows the predicate to look at the key path. Only works on .*_with_path APIs.

Fixes https://github.com/jax-ml/jax/issues/27996



Patch:
diff --git a/jax/BUILD b/jax/BUILD
index 3ec1567e2..a2d59c98e 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -1438,6 +1438,7 @@ pytype_strict_library(
     srcs = ["_src/tree.py"],
     deps = [
         ":tree_util",
+        "//jax/_src/lib",
     ],
 )
 
diff --git a/jax/_src/tree.py b/jax/_src/tree.py
index 9a3e001d9..5aaa0bf2b 100644
--- a/jax/_src/tree.py
+++ b/jax/_src/tree.py
@@ -287,7 +287,8 @@ def unflatten(treedef: tree_util.PyTreeDef,
 
 
 def flatten_with_path(
-    tree: Any, is_leaf: Callable[[Any], bool] | None = None
+    tree: Any, is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
 ) -> tuple[list[tuple[tree_util.KeyPath, Any]], tree_util.PyTreeDef]:
   """Flattens a pytree like ``tree_flatten``, but also returns each leaf's key path.
 
@@ -313,11 +314,12 @@ def flatten_with_path(
     - :func:`jax.tree.map_with_path`
     - :func:`jax.tree_util.register_pytree_with_keys`
   """
-  return tree_util.tree_flatten_with_path(tree, is_leaf)
+  return tree_util.tree_flatten_with_path(tree, is_leaf, is_leaf_takes_path)
 
 
 def leaves_with_path(
-    tree: Any, is_leaf: Callable[[Any], bool] | None = None
+    tree: Any, is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
 ) -> list[tuple[tree_util.KeyPath, Any]]:
   """Gets the leaves of a pytree like ``tree_leaves`` and returns each leaf's key path.
 
@@ -338,14 +340,15 @@ def leaves_with_path(
     - :func:`jax.tree.flatten_with_path`
     - :func:`jax.tree_util.register_pytree_with_keys`
   """
-  return tree_util.tree_leaves_with_path(tree, is_leaf)
+  return tree_util.tree_leaves_with_path(tree, is_leaf, is_leaf_takes_path)
 
 
 def map_with_path(
     f: Callable[..., Any],
     tree: Any,
     *rest: Any,
-    is_leaf: Callable[[Any], bool] | None = None,
+    is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
 ) -> Any:
   """Maps a multi-input function over pytree key path and args to produce a new pytree.
 
@@ -377,7 +380,9 @@ def map_with_path(
     - :func:`jax.tree.leaves_with_path`
     - :func:`jax.tree_util.register_pytree_with_keys`
   """
-  return tree_util.tree_map_with_path(f, tree, *rest, is_leaf=is_leaf)
+  return tree_util.tree_map_with_path(
+      f, tree, *rest, is_leaf=is_leaf, is_leaf_takes_path=is_leaf_takes_path
+  )
 
 
 def broadcast(prefix_tree: Any, full_tree: Any,
diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
index 6edbbfd62..6054b4711 100644
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -27,6 +27,7 @@ from jax._src import traceback_util
 from jax._src.lib import pytree
 from jax._src.util import safe_zip, set_module
 from jax._src.util import unzip2
+from jax._src.lib import jaxlib_extension_version
 
 
 export = set_module('jax.tree_util')
@@ -1126,34 +1127,40 @@ def register_static(cls: type[H]) -> type[H]:
 
 @export
 def tree_flatten_with_path(
-    tree: Any, is_leaf: Callable[[Any], bool] | None = None
+    tree: Any, is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
 ) -> tuple[list[tuple[KeyPath, Any]], PyTreeDef]:
   """Alias of :func:`jax.tree.flatten_with_path`."""
-  return default_registry.flatten_with_path(tree, is_leaf)
+  if jaxlib_extension_version < 351:
+    return default_registry.flatten_with_path(tree, is_leaf)
+  is_leaf_with_kp: Callable[[Any, Any], bool] | None = is_leaf
+  if not is_leaf_takes_path and is_leaf is not None:
+    is_leaf_with_kp = lambda _, x: is_leaf(x)
+  return default_registry.flatten_with_path(tree, is_leaf_with_kp)
 
 
 @export
 def tree_leaves_with_path(
-    tree: Any, is_leaf: Callable[[Any], bool] | None = None
+    tree: Any, is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
 ) -> list[tuple[KeyPath, Any]]:
   """Alias of :func:`jax.tree.leaves_with_path`."""
-  return tree_flatten_with_path(tree, is_leaf)[0]
-
-
-# generate_key_paths is not exported.
-def generate_key_paths(
-    tree: Any, is_leaf: Callable[[Any], bool] | None = None
-) -> list[tuple[KeyPath, Any]]:
-  return tree_leaves_with_path(tree, is_leaf)
-_generate_key_paths = generate_key_paths  # alias for backward compat
+  return tree_flatten_with_path(tree, is_leaf, is_leaf_takes_path)[0]
+generate_key_paths = tree_leaves_with_path
 
 
 @export
-def tree_map_with_path(f: Callable[..., Any],
-                       tree: Any, *rest: Any,
-                       is_leaf: Callable[[Any], bool] | None = None) -> Any:
+def tree_map_with_path(
+    f: Callable[..., Any],
+    tree: Any,
+    *rest: Any,
+    is_leaf: Callable[..., bool] | None = None,
+    is_leaf_takes_path: bool = False,
+) -> Any:
   """Alias of :func:`jax.tree.map_with_path`."""
-  keypath_leaves, treedef = tree_flatten_with_path(tree, is_leaf)
+  keypath_leaves, treedef = tree_flatten_with_path(
+      tree, is_leaf, is_leaf_takes_path
+  )
   keypath_leaves = list(zip(*keypath_leaves))
   all_keypath_leaves = keypath_leaves + [treedef.flatten_up_to(r) for r in rest]
   return treedef.unflatten(f(*xs) for xs in zip(*all_keypath_leaves))
diff --git a/jaxlib/_jax/pytree.pyi b/jaxlib/_jax/pytree.pyi
index ac5298c77..0cec90db6 100644
--- a/jaxlib/_jax/pytree.pyi
+++ b/jaxlib/_jax/pytree.pyi
@@ -45,7 +45,7 @@ class PyTreeRegistry:
   def flatten_with_path(
       self,
       tree: Any,
-      leaf_predicate: Callable[[Any], bool] | None = ...,
+      leaf_predicate: Callable[[Any, Any], bool] | None = ...,
   ) -> Tuple[list[Tuple[_KeyPath, Any]], PyTreeDef]: ...
   def register_node(
       self,
diff --git a/jaxlib/pytree.cc b/jaxlib/pytree.cc
index bd845c47e..6e2f7af98 100644
--- a/jaxlib/pytree.cc
+++ b/jaxlib/pytree.cc
@@ -545,16 +545,33 @@ nanobind::tuple FlattenedIndexKey::MatchArgs(nanobind::handle unused) {
   return nanobind::make_tuple("key");
 };
 
+/* static */ nb::object MakeKeyPathTuple(std::vector<nb::object>& keypath) {
+  const std::vector<nb::object>& frozen_keypath = keypath;
+  nb::object kp_tuple = nb::steal(PyTuple_New(frozen_keypath.size()));
+  for (int i = 0; i < frozen_keypath.size(); ++i) {
+    PyTuple_SET_ITEM(kp_tuple.ptr(), i,
+                      nb::object(frozen_keypath[i]).release().ptr());
+  }
+  return kp_tuple;
+}
+
 template <typename T>
-void PyTreeDef::FlattenImpl(nb::handle handle, T& leaves,
-                            const std::optional<nb::callable>& leaf_predicate,
-                            std::optional<std::vector<nb::object>>& keypath) {
+void PyTreeDef::FlattenImpl(
+    nb::handle handle, T& leaves,
+    std::optional<std::vector<nb::object>>& keypath,
+    const std::optional<nb::callable>& leaf_predicate) {
   Node node;
   const int start_num_nodes = traversal_.size();
   const int start_num_leaves = leaves.size();
   bool is_known_leaf = false;
   if (leaf_predicate) {
-    nb::object o = (*leaf_predicate)(handle);
+    nb::object o;
+    if (keypath.has_value()) {
+      auto kp_tuple = MakeKeyPathTuple(keypath.value());
+      o = (*leaf_predicate)(kp_tuple, handle);
+    } else {
+      o = (*leaf_predicate)(handle);
+    }
     // Historically we accepted "truthy" values from leaf predicates. Accept
     // None here to keep existing clients happy.
     if (o.is_none()) {
@@ -568,12 +585,7 @@ void PyTreeDef::FlattenImpl(nb::handle handle, T& leaves,
   if (is_known_leaf) {
     nb::object value = nb::borrow<nb::object>(handle);
     if (keypath.has_value()) {
-      const std::vector<nb::object>& frozen_keypath = keypath.value();
-      nb::object kp_tuple = nb::steal(PyTuple_New(frozen_keypath.size()));
-      for (int i = 0; i < frozen_keypath.size(); ++i) {
-        PyTuple_SET_ITEM(kp_tuple.ptr(), i,
-                         nb::object(frozen_keypath[i]).release().ptr());
-      }
+      auto kp_tuple = MakeKeyPathTuple(keypath.value());
       value = nb::make_tuple(std::move(kp_tuple), std::move(value));
     }
     if constexpr (std::is_same_v<T, nb::list>) {
@@ -590,7 +602,7 @@ void PyTreeDef::FlattenImpl(nb::handle handle, T& leaves,
               " in flatten; PyTree may have cyclical node references.")) {
         return;
       }
-      FlattenImpl(child, leaves, leaf_predicate, keypath);
+      FlattenImpl(child, leaves, keypath, leaf_predicate);
       Py_LeaveRecursiveCall();
     };
     switch (node.kind) {
@@ -718,12 +730,7 @@ void PyTreeDef::FlattenImpl(nb::handle handle, T& leaves,
         DCHECK(node.kind == PyTreeKind::kLeaf);
         auto value = nb::borrow<nb::object>(handle);
         if (keypath.has_value()) {
-          const std::vector<nb::object>& frozen_keypath = keypath.value();
-          nb::object kp_tuple = nb::steal(PyTuple_New(frozen_keypath.size()));
-          for (int i = 0; i < frozen_keypath.size(); ++i) {
-            PyTuple_SET_ITEM(kp_tuple.ptr(), i,
-                             nb::object(frozen_keypath[i]).release().ptr());
-          }
+          auto kp_tuple = MakeKeyPathTuple(keypath.value());
           value = nb::make_tuple(std::move(kp_tuple), std::move(value));
         }
         if constexpr (std::is_same_v<T, nb::list>) {
@@ -742,19 +749,19 @@ void PyTreeDef::Flatten(nb::handle handle,
                         absl::InlinedVector<nb::object, 2>& leaves,
                         std::optional<nb::callable> leaf_predicate) {
   std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
+  FlattenImpl(handle, leaves, keypath, leaf_predicate);
 }
 
 void PyTreeDef::Flatten(nb::handle handle, std::vector<nb::object>& leaves,
                         std::optional<nb::callable> leaf_predicate) {
   std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
+  FlattenImpl(handle, leaves, keypath, leaf_predicate);
 }
 
 void PyTreeDef::Flatten(nb::handle handle, nb::list& leaves,
                         std::optional<nb::callable> leaf_predicate) {
   std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
+  FlattenImpl(handle, leaves, keypath, leaf_predicate);
 }
 
 /*static*/ std::pair<std::vector<nb::object>, nb_class_ptr<PyTreeDef>>
@@ -766,10 +773,11 @@ PyTreeDef::Flatten(nb::handle x, nb_class_ptr<PyTreeRegistry> registry,
   return std::make_pair(std::move(leaves), std::move(def));
 }
 
-void PyTreeDef::FlattenWithPath(nb::handle handle, nanobind::list& leaves,
-                                std::optional<nb::callable> leaf_predicate) {
+void PyTreeDef::FlattenWithPath(
+    nb::handle handle, nanobind::list& leaves,
+    std::optional<nb::callable> leaf_predicate) {
   std::optional<std::vector<nb::object>> keypath = std::vector<nb::object>();
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
+  FlattenImpl(handle, leaves, keypath, leaf_predicate);
 }
 
 /*static*/ bool PyTreeDef::AllLeaves(PyTreeRegistry* registry,
diff --git a/jaxlib/pytree.h b/jaxlib/pytree.h
index 0a012d933..8f13a2613 100644
--- a/jaxlib/pytree.h
+++ b/jaxlib/pytree.h
@@ -367,9 +367,10 @@ class PyTreeDef {
       const;
 
   template <typename T>
-  void FlattenImpl(nanobind::handle handle, T& leaves,
-                   const std::optional<nanobind::callable>& leaf_predicate,
-                   std::optional<std::vector<nanobind::object>>& keypath);
+  void FlattenImpl(
+      nanobind::handle handle, T& leaves,
+      std::optional<std::vector<nanobind::object>>& keypath,
+      const std::optional<nanobind::callable>& leaf_predicate);
 
   template <typename T>
   nanobind::object UnflattenImpl(T leaves) const;
diff --git a/jaxlib/xla_client.py b/jaxlib/xla_client.py
index 2c73947fa..92657f9f2 100644
--- a/jaxlib/xla_client.py
+++ b/jaxlib/xla_client.py
@@ -43,7 +43,7 @@ ifrt_programs = _xla.ifrt_programs
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes. In JAX, reference this via jax._src.lib.jaxlib_extension_version.
-_version = 350
+_version = 351
 
 # An internal increasing version number for protecting jaxlib code against
 # ifrt changes.


Test patch:
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
index 0d92156b2..49318e41f 100644
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -26,6 +26,7 @@ from jax import flatten_util
 from jax import tree_util
 from jax._src import test_util as jtu
 from jax._src.tree_util import flatten_one_level, prefix_errors
+from jax._src.lib import jaxlib_extension_version
 import jax.numpy as jnp
 
 # Easier to read.
@@ -789,9 +790,14 @@ class TreeTest(jtu.JaxTestCase):
   def testTreeMapWithPathWithIsLeafArgument(self):
     x = ((1, 2), [3, 4, 5])
     y = (([3], jnp.array(0)), ([0], 7, [5, 6]))
-    out = tree_util.tree_map_with_path(
-        lambda kp, *xs: (kp[0].idx, *xs), x, y,
-        is_leaf=lambda n: isinstance(n, list))
+    if jaxlib_extension_version < 351:
+      out = tree_util.tree_map_with_path(
+          lambda kp, *xs: (kp[0].idx, *xs), x, y,
+          is_leaf=lambda n: isinstance(n, list))
+    else:
+      out = tree_util.tree_map_with_path(
+          lambda kp, *xs: (kp[0].idx, *xs), x, y,
+          is_leaf=lambda _, n: isinstance(n, list), is_leaf_takes_path=True)
     self.assertEqual(out, (((0, 1, [3]),
                             (0, 2, jnp.array(0))),
                            (1, [3, 4, 5], ([0], 7, [5, 6]))))
@@ -808,7 +814,13 @@ class TreeTest(jtu.JaxTestCase):
     tree1 = {'a': 1,
              'sub': [jnp.array((1, 2)), ATuple(foo=(), bar=[None])],
              'obj': AnObject2(x=EmptyTuple(), y=0, z='constantdef')}
-    flattened, _ = tree_util.tree_flatten_with_path(tree1, is_empty)
+    if jaxlib_extension_version < 351:
+      flattened, _ = tree_util.tree_flatten_with_path(tree1, is_empty)
+    else:
+      is_empty_new = lambda kp, x: is_empty(x)
+      flattened, _ = tree_util.tree_flatten_with_path(
+          tree1, is_empty_new, is_leaf_takes_path=True
+      )
     strs = [f"{tree_util.keystr(kp)}: {x}" for kp, x in flattened]
     self.assertEqual(
         strs,
@@ -822,6 +834,36 @@ class TreeTest(jtu.JaxTestCase):
         ],
     )
 
+  def testTreeFlattenWithPathWithIsLeafWithPathArgument(self):
+    if jaxlib_extension_version < 351:
+      self.skipTest("Requires jaxlib version >= 351")
+    x = ((1, 2), [3, {4: 4, 5: 5}])
+    check_max_depth = lambda kp, _: len(kp) >= 2
+    flattened, _ = tree_util.tree_flatten_with_path(
+        x, is_leaf=check_max_depth, is_leaf_takes_path=True
+    )
+    self.assertEqual(
+        flattened,
+        [
+            ((SequenceKey(0), SequenceKey(0),), 1),
+            ((SequenceKey(0), SequenceKey(1),), 2),
+            ((SequenceKey(1), SequenceKey(0),), 3),
+            ((SequenceKey(1), SequenceKey(1)), {4: 4, 5: 5}),
+        ],
+    )
+
+  def testTreeMapWithPathWithIsLeafWithPathArgument(self):
+    if jaxlib_extension_version < 351:
+      self.skipTest("Requires jaxlib version >= 351")
+    x = ((1, 2), [3, 4, 5])
+    y = (([3], jnp.array(0)), ([0], 7, [5, 6]))
+    out = tree_util.tree_map_with_path(
+        lambda kp, *xs: (kp[0].idx, *xs), x, y,
+        is_leaf=lambda kp, n: isinstance(n, list), is_leaf_takes_path=True)
+    self.assertEqual(out, (((0, 1, [3]),
+                            (0, 2, jnp.array(0))),
+                           (1, [3, 4, 5], ([0], 7, [5, 6]))))
+
   def testTreeFlattenWithPathBuiltin(self):
     x = (1, {"a": 2, "b": 3})
     flattened = tree_util.tree_flatten_with_path(x)
@@ -1522,9 +1564,16 @@ class TreeAliasTest(jtu.JaxTestCase):
   def test_tree_flatten_with_path_is_leaf(self):
     obj = [1, 2, (3, 4)]
     is_leaf = lambda x: isinstance(x, tuple)
+    if jaxlib_extension_version < 351:
+      self.assertEqual(
+          jax.tree.flatten_with_path(obj, is_leaf=is_leaf),
+          tree_util.tree_flatten_with_path(obj, is_leaf=is_leaf),
+      )
+      return
+    is_leaf = lambda kp, x: isinstance(x, tuple)
     self.assertEqual(
-        jax.tree.flatten_with_path(obj, is_leaf=is_leaf),
-        tree_util.tree_flatten_with_path(obj, is_leaf=is_leaf),
+        jax.tree.flatten_with_path(obj, is_leaf, is_leaf_takes_path=True),
+        tree_util.tree_flatten_with_path(obj, is_leaf, is_leaf_takes_path=True),
     )
 
   def test_tree_leaves_with_path(self):
@@ -1537,9 +1586,20 @@ class TreeAliasTest(jtu.JaxTestCase):
   def test_tree_leaves_with_path_is_leaf(self):
     obj = [1, 2, (3, 4)]
     is_leaf = lambda x: isinstance(x, tuple)
+    if jaxlib_extension_version < 351:
+      self.assertEqual(
+          jax.tree.leaves_with_path(obj, is_leaf=is_leaf),
+          tree_util.tree_leaves_with_path(obj, is_leaf=is_leaf),
+      )
+      return
+    is_leaf = lambda kp, x: isinstance(x, tuple)
     self.assertEqual(
-        jax.tree.leaves_with_path(obj, is_leaf=is_leaf),
-        tree_util.tree_leaves_with_path(obj, is_leaf=is_leaf),
+        jax.tree.leaves_with_path(
+            obj, is_leaf=is_leaf, is_leaf_takes_path=True
+        ),
+        tree_util.tree_leaves_with_path(
+            obj, is_leaf=is_leaf, is_leaf_takes_path=True
+        ),
     )
 
   def test_tree_map_with_path(self):
@@ -1556,9 +1616,20 @@ class TreeAliasTest(jtu.JaxTestCase):
     obj = [1, 2, (3, 4)]
     obj2 = [5, 6, (7, 8)]
     is_leaf = lambda x: isinstance(x, tuple)
+    if jaxlib_extension_version < 351:
+      self.assertEqual(
+          jax.tree.map_with_path(func, obj, obj2, is_leaf=is_leaf),
+          tree_util.tree_map_with_path(func, obj, obj2, is_leaf=is_leaf),
+      )
+      return
+    is_leaf = lambda kp, x: isinstance(x, tuple)
     self.assertEqual(
-        jax.tree.map_with_path(func, obj, obj2, is_leaf=is_leaf),
-        tree_util.tree_map_with_path(func, obj, obj2, is_leaf=is_leaf),
+        jax.tree.map_with_path(
+            func, obj, obj2, is_leaf=is_leaf, is_leaf_takes_path=True
+        ),
+        tree_util.tree_map_with_path(
+            func, obj, obj2, is_leaf=is_leaf, is_leaf_takes_path=True
+        ),
     )