`nnx.PathContains` should support regex

[thijs-vanweezel]

The current implementation of nnx.PathContains only looks for exact matches in the path. This makes it difficult to select, e.g., a set of layers which have been given a similar name. To illustrate, my use case would be to use nnx.split to separate convolutional layers and linear layers, which I respectively named "conv1", "conv2", etc. and "fc1", "fc2", etc. To my knowledge, there are also no other filters that support this.

My proposal is a very simple modification method, namely:

@dataclasses.dataclass(frozen=True)
class PathContains:
  key: Key

  def __call__(self, path: PathParts, x: tp.Any):
    pattern = re.compile(self.key)
    return any(pattern.match(segment) for segment in path)

  def __repr__(self):
    return f'PathContains({self.key!r})'

Let me know if it is worth opening a pull request for this, and whether any efficiency constraints should be taken into account.

[vfdev-5]

@thijs-vanweezel thanks for the idea! Can you give a detailed example of your model and what exactly you try to select and it is hard to get expected result with nnx.PathContains? A small reproducer would be ideal. Thanks

[thijs-vanweezel]

Hi @vfdev-5,

For a very simple example, consider the following modified LeNet-5.

import re

# LeNet-5 for 36X60 images + 3 auxiliary features
class LeNet(nnx.Module):
    def __init__(self, key):
        super().__init__()
        self.conv1 = nnx.Conv(1, 8, (4,4), rngs=key, padding="VALID")
        self.conv2 = nnx.Conv(8, 16, (4,4), rngs=key, padding="VALID")
        self.fc1 = nnx.Linear(6*12*16+3, 128, rngs=key)
        self.fc2 = nnx.Linear(128, 64, rngs=key)
        self.fc3 = nnx.Linear(64, 16, rngs=key)
    
    def __call__(self, x, z):
        x = self.conv1(x)
        x = nnx.relu(x)
        x = nnx.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = self.conv2(x)
        x = nnx.relu(x)
        x = nnx.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = jnp.reshape(x, (x.shape[0], -1))
        x = jnp.concatenate([x, z], axis=-1)
        x = self.fc1(x)
        x = nnx.relu(x)
        x = self.fc2(x)
        x = nnx.relu(x)
        x = self.fc3(x)
        return x

# Split model into fc and conv layers
struct, fc, conv, rest = nnx.split(
    model, 
    nnx.All(lambda path, x: any(re.compile(r"^fc.$").match(segment) for segment in path), nnx.PathContains("kernel")), 
    nnx.All(lambda path, x: any(re.compile(r"^conv.$").match(segment) for segment in path), nnx.PathContains("kernel")),
    ...
)

# Apply transformations with jax.tree.map
# ...

As you can see, I use lambda path, x: any(re.compile(r"^fc.$").match(segment) for segment in path), since PathContains does not support regex. Am I correct to conclude that no other filters provide this functionality?

[vfdev-5]

We can do simplier than regex, but I agree that having a built-in flexible filter class would be nice:

class PathContains:
    def __init__(self, keys: tuple[str, ...]):
        self.keys = keys

    def __call__(self, path, x):
        path_str = "".join(path)
        return all(key in path_str for key in self.keys)


struct2, fc2, conv2, rest2 = nnx.split(
    model,
    nnx.All(PathContains(["fc", "kernel"])),
    nnx.All(PathContains(["conv", "kernel"])),
    ...
)

[mohsinm-dev]

@vfdev-5 I can work on it. Also, I guess if we just change the single string from a tuple, the change is required from PathContains('head') to PathContains(('head',)) for someone already using it, or does it still support a single string? We can add the PathContainsAll class as you mentioned. If @thijs-vanweezel doesn't plan to work on it, I'd be happy to implement this approach.

[thijs-vanweezel]

I agree with @vfdev-5 that it can be done without regex. I shall open a pull request. @mohsinm-dev, feel free to contribute once it's there.

[vfdev-5]

Initially I was thinking about adding an arg to nnx.PathContains, for example, a generic path_transform or directly path_to_str (one option not both, but in the code below I show both):

# FLax NNX code
@dataclasses.dataclass(frozen=True)
class PathContains:
  key: Key
  path_transform: tp.Callable[[PathParts], Any] = lambda x: x
  path_as_str: bool

  def __call__(self, path: PathParts, x: tp.Any):
    if self.path_as_str:
      return self.key in "".join(path)
    return self.key in self.path_transform(path)

# User code:
graphdef, fc_kernels, others = nnx.split(
  model,
  nnx.All(
    nnx.Param, 

    nnx.PathContains("fc", path_to_str=True), 
    # The line below is equivalent to the line above
    # lambda p, x: "fc" in "".join(p),
     
    nnx.PathContains("kernel"),
  ),
  ...
)

I'm hesitating what to implement and I'll discuss with the team what kind of improvement if any we would like to see in Flax concerning this question.

Finally, actually the code in the question we can write simply:

struct2, fc2, conv2, rest2 = nnx.split(
    model,
    nnx.All(lambda p, x: "fc" in "".join(p), nnx.PathContains("kernel")),
    nnx.All(lambda p, x: "conv" in "".join(p), nnx.PathContains("kernel")),
    ...
)

[thijs-vanweezel]

Thanks for thinking along. I have tested the following, which raised no errors in the test files. This implementation only converts the key and the path to strings, which I think should be possible because I cannot imagine that the instances of Key used in this context would be anything different than a string. Then again, if that is true, it may be simpler to require the key argument to be a str instead of a Key. Anyhow, I will hold off from creating a pull request until further notice. Thanks!

@dataclasses.dataclass(frozen=True)
class PathContains:
  key: Key

  def __call__(self, path: PathParts, x: tp.Any):
    return str(self.key) in "".join(map(str, path))

  def __repr__(self):
    return f'PathContains({self.key!r})'