allow use to decorate the tpu function with information about whether there...

  allow use to decorate the tpu function with information about whether there are a top level while loop which is training loop. such annotation is done in estimator, so it should be transparent to most of users (except those dealing with low level api directly).

PiperOrigin-RevId: 234204932