Filtering with `nnx.PathContains('')` returns an empty state

[NKlug]

Filtering a module with nnx.PathContains('') always returns an empty state.
Consider the following example:

from flax import nnx

class CNN(nnx.Module):

  def __init__(self, *, rngs: nnx.Rngs):
    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
    ...

  def __call__(self, x):
    ...

model = CNN(rngs=nnx.Rngs(0))

graph, state, rest = nnx.split(model, nnx.PathContains(''), ...)
print(state == nnx.State({})) # True

I would argue that the empty path '' is contained in every path, much like '' in any_string is always true.
Therefore, I'd expect state to contain the full state and rest to be empty.
What do you think?

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 24.04
• Flax, jax, jaxlib versions (obtain with pip show flax jax jaxlib: Flax 0.10.2, jax 0.5.0
• Python version: 3.12

[cgarciae]

Hi @NKlug, paths in NNX are represented as a tuple[Key, ...], where Key is usually str or int. This means that PathContains will check that '' is an element of the tuple which is different. If you want to match anything you can use either ... or True as a filter. Take a look at the Filters guide.

[vfdev-5]

Let's close this issue as solved as reported behavior is expected.
@NKlug feel free to reopen the issue if we need to discuss it more.