Neel-Gupta:

Why do we need to actually use the Flax version of BigBird to run it on TPUs?


As I have said in the other topic, you don’t. Why are you using an example located in examples/flax/ if you don’t want to use FLAX?
There is really no need for the rant after, it’s completely inappropriate when people are replying to you for free on their own time.