When does recreating a model trigger nnx.jit recompilation?

[NiklasKappel]

Consider this MWE that uses the model from the MNIST tutorial:

from functools import partial

from flax import nnx
import jax.numpy as jnp


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)
        self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, 10, rngs=rngs)

    def __call__(self, x):
        x = self.avg_pool(nnx.relu(self.conv1(x)))
        x = self.avg_pool(nnx.relu(self.conv2(x)))
        x = x.reshape(x.shape[0], -1)  # flatten
        x = nnx.relu(self.linear1(x))
        x = self.linear2(x)
        return x


class Foo(nnx.Module):
    def __call__(self, x):
        return x


@nnx.jit
def foo_step(model):
    print("Compiling foo_step")
    return model(jnp.ones((1, 28, 28, 1)))


def try_CNN():
    print("Try CNN")
    print("First model")
    model = CNN(rngs=nnx.Rngs(0))
    print("First step")
    foo_step(model)  # Compiling foo_step
    print("Second step")
    foo_step(model)
    print("Second model")
    model = CNN(rngs=nnx.Rngs(0))
    print("First step")
    foo_step(model)  # Compiling foo_step, why?
    print("Second step")
    foo_step(model)


def try_Foo():
    print("Try Foo")
    print("First model")
    model = Foo()
    print("First step")
    foo_step(model)  # Compiling foo_step
    print("Second step")
    foo_step(model)
    print("Second model")
    model = Foo()
    print("First step")
    foo_step(model)  # NOT compiling foo_step
    print("Second step")
    foo_step(model)


if __name__ == "__main__":
    try_CNN()
    print()
    try_Foo()

Note that, when running try_CNN, the foo_step function is JIT compiled twice (i.e. the print side effect is triggered twice), once with the first CNN instance and once with the second, even though they are essentially the same. When running try_Foo though, the foo_step function is compiled only once. The latter is the behavior I would expect from jax.jit, considering that recreating a model does not change anything about the shape of parameter arrays etc.

What is the reason for the extra compilation happening in try_CNN?

[NiklasKappel]

Further testing reveals the culprit is the use of partial when defining the pooling layer. With the following implementation of CNN, the second compilation is avoided:

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)
        # self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
        self.avg_pool = nnx.avg_pool
        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, 10, rngs=rngs)

    def __call__(self, x):
        x = self.avg_pool(nnx.relu(self.conv1(x)), window_shape=(2, 2), strides=(2, 2))
        x = self.avg_pool(nnx.relu(self.conv2(x)), window_shape=(2, 2), strides=(2, 2))
        x = x.reshape(x.shape[0], -1)  # flatten
        x = nnx.relu(self.linear1(x))
        x = self.linear2(x)
        return x

My guess is that partial (which is a class) returns an object, and hashing this object returns its object ID, which is different for the second CNN instance and causes a cache miss with nnx.jit? If this is right, maybe its worth pointing out in the tutorial?

[DiagRisker]

Maybe the behavior is different with jax.tree_util.Partial? Have you tried?

[NiklasKappel]

I just tried and jax.tree_util.Partial avoids the additional compilation!

[cgarciae]

It's unfortunate that partial defines equality as identity. Great find @DiagRisker !