Model output various predicted result without re-init-model #1731

Description

When we load the trained model to do predict at the first time, the predict output meets our expectation, but when we try to test the identical input again without reloading (i.e., initiation) the model. The result was different and would change every time. However, if the model was initiated each time, the output could be completely identical for the same input and correct. Thus it is curious whether reloading of model is necessary before prediction of the inputs each time, because the reloading of the weight file of the model is time consuming. Any suggestion are welcome.
...

Environment information

OS: Ubuntu 20.04

$ pip freeze | grep trax
# trax==1.4.1

$ pip freeze | grep tensor
# tensorboard==2.8.0
# tensorboard-data-server==0.6.1
# tensorboard-plugin-wit==1.6.0
# tensorflow==2.8.0
# tensorflow-datasets==4.5.2
# tensorflow-gpu==2.8.0
# tensorflow-hub==0.12.0
# tensorflow-io-gcs-filesystem==0.24.0
# tensorflow-metadata==1.6.0
# tensorflow-text==2.8.1


$ pip freeze | grep jax
# jax==0.3.0

$ python -V
# Python 3.9.7