NNX MultiHeadDotProductAttention

[PVirie]

I've been using flax for a while and it's the best neural network module available as of now, to say the least.

My question is short:
Could you share the road map for an implementation for MultiHeadDotProductAttention module in NNX?
I wish to know when should I decide to upgrade from linen one.

Thank you in advance.

[PVirie]

Nvm I saw it in the nnx source.
If you have time, please update the document for how to use. An example would also be nice :)

[IvyZX]

Yes we simplified the name to MultiHeadAttention. Please check out our API reference for how to use it.