Linen modules get detailed trace annotations by default. Can NNX modules also enable them?

[NiklasKappel]

The following MWE uses jax.profiler.trace to profile the inference time of the CNN modules (with the same architecture) from the Flax Linen and Flax NNX documentations, respectively.

MWE

from timeit import timeit

import jax
from flax import linen as nn  # Linen API
from flax import nnx  # The Flax NNX API.


class LinenCNN(nn.Module):
    """A simple CNN model."""

    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = nn.Dense(features=256)(x)
        x = nn.relu(x)
        x = nn.Dense(features=10)(x)
        return x


@jax.tree_util.Partial(jax.jit, static_argnums=0)
def linen_inference_step(model_apply, params, batch):
    print("Compiling linen inference step...")
    return model_apply(params, batch)


def profile_linen_inference(model, params, batch, num_iterations, log_dir):
    # Compile.
    linen_inference_step(model.apply, params, batch)

    def time_me():
        linen_inference_step(model.apply, params, batch).block_until_ready()

    with jax.profiler.trace(log_dir):
        time = timeit(time_me, number=num_iterations)

    print(f"Time: {time}")


class NNXCNN(nnx.Module):
    """A simple CNN model."""

    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 = jax.tree_util.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


@nnx.jit
def nnx_inference_step(model, batch):
    print("Compiling nnx inference step...")
    return model(batch)


def profile_nnx_inference(model, batch, num_iterations, log_dir):
    # Compile.
    nnx_inference_step(model, batch)

    def time_me():
        nnx_inference_step(model, batch).block_until_ready()

    with jax.profiler.trace(log_dir):
        time = timeit(time_me, number=num_iterations)

    print(f"Time: {time}")


def main():
    batch = jax.random.normal(key=jax.random.PRNGKey(0), shape=(32, 28, 28, 1))

    linen_cnn = LinenCNN()
    linen_params = linen_cnn.init(jax.random.PRNGKey(0), batch)
    profile_linen_inference(linen_cnn, linen_params, batch, 100, "linen_tb_trace/")

    nnx_cnn = NNXCNN(rngs=nnx.Rngs(0))
    profile_nnx_inference(nnx_cnn, batch, 100, "nnx_tb_trace/")


if __name__ == "__main__":
    main()

When viewing the trace for the Linen model with tensorboard, the relevant section displays the names of the called modules under the "Framework Name Scope" section (here "LinenCNN", "Conv_*" and "Dense_*"):

In the trace for the NNX model, this information is missing:

Is there a built-in way to enable the same kind of annotations for NNX modules? The next best thing I can think of would be to decorate the __call__ methods of all my NNX modules with jax.named_scope.