How to get word embeddings for Flax model?

Flax/JAX Projects
1

For example model.bert.embeddings.weight gives us embeddings for my model which is:

model = TFBertForQuestionAnswering.from_pretrained("deepset/bert-base-cased-squad2", from_pt=True)

Similarly I am working on Flax model and it doesn’t has that method: my flax model is something like this:

model = FlaxBertForQuestionAnswering.from_pretrained('mrm8488/bert-multi-cased-finedtuned-xquad-tydiqa-goldp')
2

The word embeddings for flax model is: model.params[“transformer”][“wte”][“embedding”]