Help porting Gemma 3n

[leighgable]

Hello all,
I'm trying to use penzai to port over the Gemma 3n architecture. I started with the feedforward block, and before I go farther, I wanted to check in, and see if there's stuff about my approach that I should change now, before moving on the the Attention block and the rest of the model. If anyone has a moment to glance at this and make suggestions, I'd be grateful.

class FFN_Block(pz.nn.Sequential):
    """ **side_inputs? """ ```

"""And then I am using a mask to change width of the matryoshka weights x activations:"""

``` @pz.pytree_dataclass()
class FFNMask(pz.nn.Layer):
    def __call__(self, x, **side_inputs):
        width = side_inputs.get("ffn_width", 8192)
        mask = jnp.arange(x.data_array.shape[-1]) < width
        mask = pz.nx.wrap(mask).tag("intermediate")
        return x * mask 

def build_gemma3n_ffn_block(name: str, block_num: int, config: dict):
    # Projections
    act_fn = functools.partial(jax.nn.gelu, approximate=True)
    gate_proj = pz.nn.Linear.from_config(
        name=f"{name}/{block_num}/gate_proj",
        init_base_rng=jax.random.key(0),
        input_axes={"embedding": config["hidden_size"]},
        output_axes={"intermediate": config["intermediate_size"][block_num]}
    )
    up_proj = pz.nn.Linear.from_config(
        name=f"{name}/{block_num}/up_proj",
        init_base_rng=jax.random.key(0),
        input_axes={"embedding": config["hidden_size"]},
        output_axes={"intermediate": config["intermediate_size"][block_num]}
    )
    down_proj = pz.nn.Linear.from_config(
        name=f"{name}/{block_num}/down_proj",
        init_base_rng=jax.random.key(0),
        input_axes={"intermediate": config["intermediate_size"][block_num]},
        output_axes={"embedding": config["hidden_size"]}
    )

    return FFN_Block([
        # parallel projections for Gate and Up paths
        pz.nn.BranchAndMultiplyTogether([
            pz.nn.Sequential([gate_proj, # Path A
                              FFNMask(),
                             ]),
            pz.nn.Sequential([up_proj,   # Path B
                              FFNMask(),
                              pz.nn.Elementwise(act_fn)
                             ])
                    
        ]),
        
        # project back to hidden size
        down_proj
    ])