Reuse the mechanism of py_utils to have empty name for EMA variable_scope.

Currently, EMA variable scope name is the same to tf.get_variable_scope().
The implementation relies on who is caller. If the caller's variable_scope has
a name, the code assumption is broken.

py_utils.CreateVariable has same requirement, which removes variable_scope
prefix (e.g. train, evaler_cpu). So py_utils.GetLingvoVariableCreator() is used
when creating tf.variable, which works no matter graph or eager mode.

Let EMA variable scope reuse py_utils.GetLingvoVariableCreator(). EMA variable
is just variable, so it doesn't need special treatment.

In addition, remove unused _EMAVariableScope().

PiperOrigin-RevId: 487879527

Author: ds-hwang

lingvo/core/base_layer.py

@@ -893,23 +893,6 @@ def _SelfVariableScope(self, params=None, enter_name_scope=True):
             tf.name_scope(self._self_variable_scope.original_name_scope))
       yield stack
 
-  @contextlib.contextmanager
-  def _EMAVariableScope(self):
-    """Internal. Used to ensure that the EMA variable names are reused."""
-    if not hasattr(self, '_ema_variable_scope'):
-      # Use empty name here so no prefix is added to the EMA variable names.
-      # This ensures compatibility between Graph and Eager checkpoint names.
-      # https://www.tensorflow.org/api_docs/python/tf/compat/v1/variable_scope
-      with tf.variable_scope('') as scope:
-        self._ema_variable_scope = scope
-    with contextlib.ExitStack() as stack:
-      # Entering a variable scope so that the EMA variables can reuse the names
-      # in separate trainer runs. This turned out to be useful in the case of
-      # model fine tuning, when the same model is run multiple times in the same
-      # Eager context.
-      stack.enter_context(tf.variable_scope(self._ema_variable_scope))
-      yield stack
-
   def _child_variable_scope_override(self):
     """Override the variable scope for individual children.
 

lingvo/core/base_model.py

@@ -15,7 +15,6 @@
 """Base model."""
 
 import collections
-import contextlib
 import dataclasses
 import re
 from typing import Dict, Union
@@ -850,11 +849,12 @@ def ApplyExponentialMovingAverage(self, ema):
     self._graphs_applied_ema.add(graph)
 
     tf.logging.info('ApplyExponentialMovingAverage on %s', self)
-    with contextlib.ExitStack() as context_stack:
-      if py_utils.IsEagerMode():
-        context_stack.enter_context(self._EMAVariableScope())
+
+    def ApplyEma():
       with tf.name_scope('moving_average'):
         self._post_train_ops.append(ema.apply(all_vars))
+    # Use empty name here so no prefix is added to the EMA variable names.
+    py_utils.GetLingvoVariableCreator('', '')(ApplyEma)
 
   # TODO(blee): Rename Decode->DecodeWithDefaultTheta, DecodeWithTheta->Decode.
   def Decode(self, input_batch):