Set the name when creating layers.PositionEmbedding object.

PiperOrigin-RevId: 307500045

official/nlp/modeling/networks/albert_transformer_encoder.py

@@ -122,7 +122,8 @@ class AlbertTransformerEncoder(network.Network):
     self._position_embedding_layer = layers.PositionEmbedding(
         initializer=initializer,
         use_dynamic_slicing=True,
-        max_sequence_length=max_sequence_length)
+        max_sequence_length=max_sequence_length,
+        name='position_embedding')
     position_embeddings = self._position_embedding_layer(word_embeddings)
 
     type_embeddings = (

official/nlp/modeling/networks/encoder_scaffold.py

@@ -146,7 +146,8 @@ class EncoderScaffold(network.Network):
       self._position_embedding_layer = layers.PositionEmbedding(
           initializer=embedding_cfg['initializer'],
           use_dynamic_slicing=True,
-          max_sequence_length=embedding_cfg['max_seq_length'])
+          max_sequence_length=embedding_cfg['max_seq_length'],
+          name='position_embedding')
       position_embeddings = self._position_embedding_layer(word_embeddings)
 
       type_embeddings = (

official/nlp/modeling/networks/transformer_encoder.py

@@ -118,7 +118,8 @@ class TransformerEncoder(network.Network):
     self._position_embedding_layer = layers.PositionEmbedding(
         initializer=initializer,
         use_dynamic_slicing=True,
-        max_sequence_length=max_sequence_length)
+        max_sequence_length=max_sequence_length,
+        name='position_embedding')
     position_embeddings = self._position_embedding_layer(word_embeddings)
 
     type_embeddings = (