How to pass a Module when initializing another Module

[Zirunis]

I feel a little stupid, but I can't figure it out and I can't find anything online either (which most likely has to do with me not knowing how to properly describe what I want to do to a search engine).

Basically I just want to pass a nnx.Module as a parameter when initializing my Model.

What I have (only bottom 5 lines should matter):

class AbstractRelativeBias(abc.ABC, nnx.Module):
    @abc.abstractmethod
    def __init__(self, n_heads: int, *args, rngs: nnx.Rngs, **kwargs):
        raise NotImplementedError

    @abc.abstractmethod
    def __call__(self, t_q: Real[Array, "b l"], t_k: Real[Array, "b t"]) -> Float[Array, "b h l t"]:
        raise NotImplementedError


class RelativePositionBias(AbstractRelativeBias):
    def __init__(self, n_heads: int, max_dist=256, *, rngs: nnx.Rngs):
        self.max_dist = max_dist
        self.bias = nnx.Embed(2 * max_dist + 1, n_heads, rngs=rngs)

    def __call__(self, t_q: Integer[Array, "b l"], t_k: Integer[Array, "b t"]) -> Float[Array, "b h l t"]:
        rel_pos = t_q[:, :, None] - t_k[:, None, :]  # shape: (b, l, t)
        # clip and shift into positive for embedding lookup
        rel_pos = rel_pos.clip(-self.max_dist, self.max_dist) + self.max_dist
        bias = self.bias(rel_pos)  # lookup embeddings: (b, l, t, n_heads)
        return rearrange(bias, 'b l t h -> b h l t')


class TransformerEncoder(nnx.Module):
    def __init__(
        self,
        num_layers: int,
        d_model: int,
        n_heads: int,
        d_ff: int,
        *,
        dropout: float = 0.0,
        rel_pos_bias_cls: Type[AbstractRelativeBias] = RelativePositionBias,  # <--- I want to be able to pass different Module classes
        rngs: nnx.Rngs
    ):
        super().__init__()

        rel_pos_bias = rel_pos_bias_cls(n_heads, rngs=rngs)  # <--- That are instantiated here

        @nnx.split_rngs(splits=num_layers)
        @nnx.vmap(in_axes=(0,), out_axes=0)
        def create_block(rngs: nnx.Rngs):
            return TransformerBlock(
                d_model,
                n_heads,
                d_ff,
                dropout=dropout,
                rel_pos_bias_module=rel_pos_bias,
                rngs=rngs
            )

        self.blocks = create_block(rngs)
        self.final_ln = nnx.LayerNorm(d_model, rngs=rngs)

What I'm getting:

Error in call to target 'src.models.transformer.TransformerEncoder':
ValueError('Cannot extract graph node from different trace level, got Param( # 2,052 (8.2 KB)\n  value=Array([[ 0.21615325,  0.2936319 , -0.57083714, -0.18689953],\n         [-0.09955087, -0.8635547 , -0.91651356, -0.23084189],\n         [-0.01597754, -0.88869476,  0.70773613,  0.07927866],\n         ...,\n         [-0.58450574, -0.591149  ,  0.04651174, -0.40675947],\n         [ 0.29281357, -0.4851477 ,  0.1312826 , -0.5893977 ],\n         [-0.21419519, -0.7562914 ,  0.27110362, -0.3196758 ]],      dtype=float32)\n)')
full_key: encoder.model

Maybe important for context: My train step is jitted using jax.jit (with nnx.merge and nnx.split to make it work) and tracing (and running) worked fine before I added this functionality.

If this is the completely wrong approach, just let me know, I'd actually be happy to learn the cleanest way to do this.

[Zirunis]

Ah, I figured out the problem. I was passing the same instance to each TransformerBlock right below. Will edit the original post to show.

For anyone with the same problem, this is the correct pattern if each layer needs an instance (kind of obvious in hindsight):

class TransformerEncoder(nnx.Module):
    def __init__(
        self,
        num_layers: int,
        d_model: int,
        n_heads: int,
        d_ff: int,
        *,
        dropout: float = 0.0,
        rel_pos_bias_cls: Type[AbstractRelativeBias] = RelativePositionBias,
        rngs
    ):
        super().__init__()

        @nnx.split_rngs(splits=num_layers)
        @nnx.vmap(in_axes=(0,), out_axes=0)
        def create_block(rngs: nnx.Rngs):
            rel_pos_bias = rel_pos_bias_cls(n_heads, rngs=rngs)  # <-- give each block it's own bias
            return TransformerBlock(
                d_model,
                n_heads,
                d_ff,
                dropout=dropout,
                rel_pos_bias_module=rel_pos_bias,
                rngs=rngs
            )

        self.blocks = create_block(rngs)
        self.final_ln = nnx.LayerNorm(d_model, rngs=rngs)

If you need the parameters to be shared, this ain't for you. Not sure how that is done, but I'll have to figure it out soon.