tensorflow.train import crashes in bert.py

[friesel]

Description

In models/bert.py the line
from tensorflow.train import load_checkpoint
crashes with
Traceback (most recent call last):
File "math_trax.py", line 19, in
import trax
File "/root/.local/lib/python3.6/site-packages/trax/init.py", line 19, in
from trax import lr_schedules as lr
File "/root/.local/lib/python3.6/site-packages/trax/lr_schedules.py", line 37, in
from trax import models as trax_models
File "/root/.local/lib/python3.6/site-packages/trax/models/init.py", line 32, in
from trax.models.research import bert
File "/root/.local/lib/python3.6/site-packages/trax/models/research/bert.py", line 20, in
from tensorflow.train import load_checkpoint
ModuleNotFoundError: No module named 'tensorflow.train'

If we change that (in bert.py) to
#from tensorflow.train import load_checkpoint
from tensorflow_core._api.v2.train import load_checkpoint

the import works.

My environment or setup? Sorry if so. Tried to exclude that but to no avail.

...

Environment information

OS: ubuntu 18.04 in a docker container

mesh-tensorflow==0.1.11
tensor2tensor==1.15.4
tensorboard==2.1.0
tensorflow==2.1.0
tensorflow-datasets==2.1.0
tensorflow-estimator==2.1.0
tensorflow-gan==2.0.0
tensorflow-hub==0.7.0
tensorflow-metadata==0.21.1
tensorflow-probability==0.7.0
# your output here

$ pip freeze | grep jax
jax==0.1.59
jaxlib==0.1.39
(I know I need to upgrade but I deem unrelated)

$ python3 -V
Python 3.6.9

For bugs: reproduction and error logs

# Steps to reproduce:

import trax

Error logs:

Traceback (most recent call last):
File "math_trax.py", line 19, in
import trax
File "/root/.local/lib/python3.6/site-packages/trax/init.py", line 19, in
from trax import lr_schedules as lr
File "/root/.local/lib/python3.6/site-packages/trax/lr_schedules.py", line 37, in
from trax import models as trax_models
File "/root/.local/lib/python3.6/site-packages/trax/models/init.py", line 32, in
from trax.models.research import bert
File "/root/.local/lib/python3.6/site-packages/trax/models/research/bert.py", line 20, in
from tensorflow.train import load_checkpoint
ModuleNotFoundError: No module named 'tensorflow.train'

[lukaszkaiser]

This is corrected at head, thanks for reporting! (And please reopen if there are more problems of course!)