[MTF] Maybe context.length_dim doesn't exist outside context.train?

PiperOrigin-RevId: 361157130

Author: afrozenator

mesh_tensorflow/transformer/evolved_transformer.py

@@ -228,7 +228,10 @@ def _pad_channels_dim(tensor, size):
 
 
 def _dropout(x, context, dropout_rate):
-  return mtf.dropout(
-      x, context.train,
-      rate=dropout_rate,
-      noise_shape=mtf.Shape(context.batch_dims + x.shape.dims[-1:]))
+  if context.train and dropout_rate > 0:
+    return mtf.dropout(
+        x, context.train,
+        rate=dropout_rate,
+        noise_shape=mtf.Shape(context.batch_dims + x.shape.dims[-1:]))
+  else:
+    return x

mesh_tensorflow/transformer/fixup_layers.py

@@ -316,9 +316,10 @@ def call(self, context, x, losses=None):
         name="wi",
         kernel_initializer=self.upproject_initializer,
         expert_dims=context.model.ensemble_dims)
-    h = mtf.dropout(
-        h, context.train, 1.0 - self.dropout_rate,
-        noise_shape=h.shape - context.length_dim)
+    if context.train and self.dropout_rate != 0.0:
+      h = mtf.dropout(
+          h, context.train, 1.0 - self.dropout_rate,
+          noise_shape=h.shape - context.length_dim)
     shift = get_single_scalar_bias(x, "shift")
     h_res = mtf.add(h, shift)
     h = mtf.reshape(h_res, h.shape)

mesh_tensorflow/transformer/transformer.py

@@ -587,9 +587,12 @@ def sublayer_residual(x, layer_stack, context):
 @gin.configurable
 def sublayer_dropout(x, layer_stack, context, dropout_rate=0.0):
   del layer_stack
-  return mtf.dropout(
-      x, context.train, rate=dropout_rate,
-      noise_shape=mtf.Shape(context.batch_dims + [context.model.model_dim]))
+  if context.train and dropout_rate > 0:
+    return mtf.dropout(
+        x, context.train, rate=dropout_rate,
+        noise_shape=mtf.Shape(context.batch_dims + [context.model.model_dim]))
+  else:
+    return x
 
 
 @gin.configurable

mesh_tensorflow/transformer/transformer_layers.py

@@ -64,8 +64,9 @@ def call(self, context, x, losses=None):
                                  variable_dtype=context.variable_dtype,
                                  name="wi",
                                  expert_dims=context.model.ensemble_dims)
-    h = mtf.dropout(h, context.train, 1.0 - self.dropout_rate,
-                    noise_shape=h.shape - context.length_dim)
+    if context.train and self.dropout_rate != 0.0:
+      h = mtf.dropout(h, context.train, keep_prob=1.0 - self.dropout_rate,
+                      noise_shape=h.shape - context.length_dim)
     return mtf.layers.dense(h, io_channels,
                             use_bias=self.use_bias,
                             activation=None,

mesh_tensorflow/transformer/universal_transformer.py

@@ -697,10 +697,13 @@ def call(self, context, x):
     return x
 
   def _dropout(self, context, x):
-    return mtf.dropout(
-        x, context.train,
-        rate=self._dropout_rate,
-        noise_shape=mtf.Shape(context.batch_dims + [context.model.model_dim]))
+    if context.train and self._dropout_rate > 0:
+      return mtf.dropout(
+          x, context.train,
+          rate=self._dropout_rate,
+          noise_shape=mtf.Shape(context.batch_dims + [context.model.model_dim]))
+    else:
+      return x
 
   def _layer_norm(self, context, x, name=None):
     """Layer normalization.

mesh_tensorflow/transformer/vocab_embeddings.py

@@ -509,11 +509,13 @@ def _sigmoid_tree(self, tensor):
     ], self._gates_dim.name)
 
   def _dropout(self, tensor, context):
-    return mtf.dropout(
-        tensor,
-        context.train,
-        1.0 - self._dropout_rate,
-        noise_shape=tensor.shape - context.length_dim)
+    if context.train and self._dropout_rate != 0.0:
+      return mtf.dropout(
+          tensor,
+          context.train,
+          1.0 - self._dropout_rate,
+          noise_shape=tensor.shape - context.length_dim)
+    return tensor
 
   def _rearrange_sentinels(self, logits):
     """Reorder along the vocab dim so the last few tokens don't share gates."""