TensorFlow CI fixes
I made a lot of changes to the TF tests, and this exposed a few issues. This PR fixes all the exposed issues, so hopefully after this the only remaining CI issues should be related to generation or the `SharedEmbeddings` refactor.
_The documentation is not available anymore as the PR was closed or merged._