Repository: patrick-kidger/equinox (language: python)

Issue:
Add `is_leaf` argument to `filter` and `partition`
Same behaviour as in `jax.tree_map` etc.


Discussion / hints:

Fix patch:
diff --git a/equinox/filters.py b/equinox/filters.py
index 6bd9c0a6..50574c3f 100644
--- a/equinox/filters.py
+++ b/equinox/filters.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Callable, Optional
 
 import jax
 import jax.numpy as jnp
@@ -50,13 +50,18 @@ def is_inexact_array_like(element: Any) -> bool:
 #
 
 
-def _make_filter_tree(mask: BoolAxisSpec, arg: Any) -> ResolvedBoolAxisSpec:
-    if isinstance(mask, bool):
-        return jax.tree_map(lambda _: mask, arg)
-    elif callable(mask):
-        return jax.tree_map(mask, arg)
-    else:
-        raise ValueError("`filter_spec` must consist of booleans and callables only.")
+def _make_filter_tree(is_leaf):
+    def _filter_tree(mask: BoolAxisSpec, arg: Any) -> ResolvedBoolAxisSpec:
+        if isinstance(mask, bool):
+            return jax.tree_map(lambda _: mask, arg, is_leaf=is_leaf)
+        elif callable(mask):
+            return jax.tree_map(mask, arg, is_leaf=is_leaf)
+        else:
+            raise ValueError(
+                "`filter_spec` must consist of booleans and callables only."
+            )
+
+    return _filter_tree
 
 
 def filter(
@@ -64,6 +69,7 @@ def filter(
     filter_spec: PyTree[BoolAxisSpec],
     inverse: bool = False,
     replace: Any = None,
+    is_leaf: Optional[Callable[[Any], bool]] = None,
 ) -> PyTree:
     """
     Filters out the leaves of a PyTree not satisfying a condition. Those not satisfying
@@ -81,6 +87,11 @@ def filter(
     - `inverse` switches the truthy/falsey behaviour: falsey results are kept and
         truthy results are replaced.
     - `replace` is what to replace any falsey leaves with. Defaults to `None`.
+    - `is_leaf`: Optional function called at each node of the PyTree. It should return
+        a boolean. `True` indicates that the whole subtree should be treated as leaf;
+        `False` indicates that the subtree should be traversed as a PyTree. This is
+        mostly useful for evaluating a callable `filter_spec` on a node instead of a
+        leaf.
 
     **Returns:**
 
@@ -94,14 +105,17 @@ def filter(
     """
 
     inverse = bool(inverse)  # just in case, to make the != trick below work reliably
-    filter_tree = jax.tree_map(_make_filter_tree, filter_spec, pytree)
+    filter_tree = jax.tree_map(_make_filter_tree(is_leaf), filter_spec, pytree)
     return jax.tree_map(
         lambda mask, x: x if bool(mask) != inverse else replace, filter_tree, pytree
     )
 
 
 def partition(
-    pytree: PyTree, filter_spec: PyTree[BoolAxisSpec], replace: Any = None
+    pytree: PyTree,
+    filter_spec: PyTree[BoolAxisSpec],
+    replace: Any = None,
+    is_leaf: Optional[Callable[[Any], bool]] = None,
 ) -> PyTree:
     """Equivalent to `filter(...), filter(..., inverse=True)`, but slightly more
     efficient.
@@ -111,7 +125,7 @@ def partition(
         See also [`equinox.combine`][] to reconstitute the PyTree again.
     """
 
-    filter_tree = jax.tree_map(_make_filter_tree, filter_spec, pytree)
+    filter_tree = jax.tree_map(_make_filter_tree(is_leaf), filter_spec, pytree)
     left = jax.tree_map(lambda mask, x: x if mask else replace, filter_tree, pytree)
     right = jax.tree_map(lambda mask, x: replace if mask else x, filter_tree, pytree)
     return left, right


Test patch:
diff --git a/tests/test_filters.py b/tests/test_filters.py
index e341cc06..e64209ef 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -144,3 +146,23 @@ def test_partition_and_combine(getkey):
 def test_partition_subtree():
     a, b = eqx.partition([(1,), 2], [True, False])
     eqx.combine(a, b)
+
+
+def test_is_leaf():
+    class M(eqx.Module):
+        value: Any
+
+    def is_m(x):
+        return isinstance(x, M)
+
+    def filter_spec(x):
+        if is_m(x):
+            return x.value == 1
+        return True
+
+    pytree = [M(1), M(2), 3]
+    out = eqx.filter(pytree, filter_spec, is_leaf=is_m)
+    assert out == [M(1), None, 3]
+    out1, out2 = eqx.partition(pytree, filter_spec, is_leaf=is_m)
+    assert out1 == [M(1), None, 3]
+    assert out2 == [None, M(2), None]