[nnx] add PathContains Filter

flax/nnx/__init__.py

@@ -25,6 +25,7 @@
 from .nnx import compat as compat
 from .nnx import traversals as traversals
 from .nnx.filterlib import All as All
+from .nnx.filterlib import PathContains as PathContains
 from .nnx.filterlib import Not as Not
 from .nnx.graph import GraphDef as GraphDef
 from .nnx.graph import GraphState as GraphState

flax/nnx/nnx/filterlib.py

@@ -14,7 +14,7 @@
 
 import builtins
 import dataclasses
-from flax.typing import PathParts
+from flax.typing import Key, PathParts
 import typing as tp
 
 if tp.TYPE_CHECKING:
@@ -66,6 +66,13 @@ class WithTag:
   def __call__(self, path: PathParts, x: tp.Any):
     return isinstance(x, _HasTag) and x.tag == self.tag
 
+@dataclasses.dataclass
+class PathContains:
+  key: Key
+
+  def __call__(self, path: PathParts, x: tp.Any):
+    return self.key in path
+
 
 @dataclasses.dataclass
 class OfType:

flax/nnx/tests/filters_test.py

@@ -0,0 +1,33 @@
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from absl.testing import absltest
+
+from flax import nnx
+
+
+class TestFilters(absltest.TestCase):
+  def test_path_contains(self):
+    class Model(nnx.Module):
+      def __init__(self, rngs):
+        self.backbone = nnx.Linear(2, 3, rngs=rngs)
+        self.head = nnx.Linear(3, 10, rngs=rngs)
+
+    model = Model(nnx.Rngs(0))
+
+    head_state = nnx.state(model, nnx.PathContains('head'))
+
+    self.assertIn('head', head_state)
+    self.assertNotIn('backbone', head_state)