Adding transformer encoder and decoder layers to flax source as in pytorch

[coder0143]

The pytorch source consists of implementations of wrappers for transformer modules.

SRC: https://github.com/pytorch/pytorch/blob/v2.9.1/torch/nn/modules/transformer.py#L966

I want to add such implementation for ease of use / ux. I will make a new file: flax/flax/nnx/nn/transformer.py which will contain the following modules:

• TransformerEncoderLayer
• TransformerEncoder
• TransformerDecoderLayer
• TransformerDecoder
• Transformer

I will keep it consistent with nnx.Linear and nnx.MultiHeadAttention modules and update the docs too, if needed, I can implement custom separate attentions such as MHSA(for full) and GQA(with kv-cache) based on review. Can I do a PR? @cgarciae @vfdev-5

[samanklesaria]

The comment for PyTorch's TransformerDecoder says the following:

The intent of this layer is as a reference implementation for foundational understanding
and thus it contains only limited features relative to newer Transformer architectures.
Given the fast pace of innovation in transformer-like architectures, we recommend
exploring this tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>_
to build efficient layers from building blocks in core or using higher
level libraries from the PyTorch Ecosystem <https://landscape.pytorch.org/>_.

That is to say, these layers in PyTorch are primarily for the sake of pedagogy. If the official recommendations from PyTorch are to "build efficient layers from building blocks in core" instead, it seems like that's the direction we should steer flax users as well.

It seems like GQA is already possible in jax if the number of query heads is different from the number of key/value heads (see jax.nn.dot_product_attention). But flax.nnx.dot_product_attention doesn't inherit this functionality. I'll open a fix for that.

[coder0143]

Can I do a PR for it, can you review it for me @samanklesaria

[samanklesaria]

@coder0143 looks like @ayulockedin beat you to it with #5180