Rewrites BERT in Flax to the new Linen API
One small caveat: we renamed `param` in Module to `params`, so we may have to update this once we created a new pypi wheel, but I will make sure that this is the case.

Otherwise it seems good to go!