Prepare for update to python3.10

PiperOrigin-RevId: 487889252

Author: lingvo-bot

lingvo/core/checkpointer_eager_test.py

@@ -116,7 +116,7 @@ def testEagerEMACheckpointCompatibility(self):
     eager_v2_logdir = os.path.join(self.get_temp_dir(), 'eager_v2')
     mdl = cfg.Instantiate()
 
-    @tf.function
+    @tf.function(autograph=False)
     def _Update():
       with py_utils.GradientTape(persistent=True):
         mdl.ConstructFPropBPropGraph()
@@ -207,7 +207,7 @@ def testEagerCheckpointConsumptionCheck(self):
     eager_v2_logdir = os.path.join(self.get_temp_dir(), 'eager_v2')
     mdl = cfg.Instantiate()
 
-    @tf.function
+    @tf.function(autograph=False)
     def _Update():
       with py_utils.GradientTape(persistent=True):
         mdl.ConstructFPropBPropGraph()

lingvo/core/layers_test.py

@@ -5730,6 +5730,9 @@ class DeterministicDropoutTest(test_utils.TestCase, parameterized.TestCase):
   def testDeterministicDropoutLayer(self):
 
     with self.session(graph=tf.Graph()):
+      # Explicitly reset step seed as `self.session(graph=tf.Graph())` does not
+      # reset the default graph id in Eager mode.
+      py_utils.ResetStepSeed()
       tf.random.set_seed(12345)
       params = layers.DeterministicDropoutLayer.Params().Set(
           name='drop', keep_prob=0.7)
@@ -5758,6 +5761,9 @@ def testDeterministicDropoutLayer(self):
     with contextlib.ExitStack() as context_stack:
       g = _ResetTfStatus(self, context_stack)
     with self.session(graph=g):
+      # Explicitly reset step seed as `self.session(graph=tf.Graph())` does not
+      # reset the default graph id in Eager mode.
+      py_utils.ResetStepSeed()
       tf.random.set_seed(12345)
       params = layers.DeterministicDropoutLayer.Params().Set(
           name='drop', keep_prob=0.7)

lingvo/core/optimizer_eager_test.py

@@ -151,7 +151,7 @@ def _Apply1(proj_layer, opt):
     # Get `snapshots` of the variables
     vars2 = [v.read_value() for v in proj_layer.vars.Flatten()]
 
-    @tf.function
+    @tf.function(autograph=False)
     def _Apply2(proj_layer, opt):
       inputs1 = np_input1
       output1 = proj_layer.FPropDefaultTheta(inputs1, in_padding1)