Does flax have an implementation of a full decoder-only model?

[epignatelli]

As per title, does flax have an implementation of a full decoder-only model, detached from its use in NPL?
I mean a generic implementation that can be used, for example, in RL.

Till now, I have found:

• The elementary mechanisms (MHA, ...)
• Pretrained models (huggingface, ...)

but I have not found the implementation of a full transformer.

Can anybody point me to one?

[mmhamdy]

Hi, AFAIK Flax does not provide implementations for ~complex transformer components such as TransformerEncoder and TransformerDecoder and just focuses on primitive building blocks. For more information, check the discussion in this PR #1480

[davisyoshida]

@epignatelli In general you can just take one of the model classes from transformers, and not use the pretrained weights.

[epignatelli]

Thanks @davisyoshida, The problem with those is that they have a lot of bells and whistles that, if not needed, make everything very hard to read, maintain and debug -- that's why I was looking for a plain implementation.