Add reduce_dims to avoid warning.

PiperOrigin-RevId: 281172010

Author: adarob

mesh_tensorflow/transformer/moe.py

@@ -368,12 +368,13 @@ def transformer_moe_layer_v1(
   # Now feed the expert inputs through the experts.
   h = mtf.layers.dense(
       expert_inputs, hidden_dim, expert_dims=[experts_dim],
+      reduced_dims=expert_inputs.shape.dims[-1:],
       activation=activation, use_bias=False,
       variable_dtype=variable_dtype, name="wi")
 
   expert_output = mtf.layers.dense(
       h, output_dim, expert_dims=[experts_dim], use_bias=False,
-      variable_dtype=variable_dtype,
+      reduced_dims=h.shape.dims[-1:], variable_dtype=variable_dtype,
       name="wo")
 
   expert_output = mtf.reshape(
@@ -630,10 +631,12 @@ def transformer_moe_layer_v2(
 
   hidden_output = mtf.layers.dense(
       expert_inputs_y, hidden_dim, expert_dims=[y0, x1],
+      reduced_dims=expert_inputs_y.shape.dims[-1:],
       activation=mtf.relu, use_bias=False, variable_dtype=variable_dtype,
       name="wi")
   expert_output = mtf.layers.dense(
       hidden_output, output_dim, expert_dims=[y0, x1],
+      reduced_dims=hidden_output.shape.dims[-1:],
       use_bias=False, variable_dtype=variable_dtype,
       name="wo")
 

mesh_tensorflow/transformer/transformer.py

@@ -591,6 +591,7 @@ def _call_internal(self, context, inputs, targets=None):
       logits = mtf.layers.dense(
           x, self.output_vocab_dim, use_bias=False,
           variable_dtype=context.variable_dtype,
+          reduced_dims=x.shape.dims[-1:],
           name="logits")
     if targets is not None and context.losses is not None:
       context.losses.append(

mesh_tensorflow/transformer/transformer_layers.py

@@ -56,12 +56,14 @@ def call(self, context, x, losses=None):
     h = mtf.layers.dense(x, hidden_channels,
                          use_bias=False, activation=mtf.relu,
                          variable_dtype=context.variable_dtype,
+                         reduced_dims=x.shape.dims[-1:],
                          name="wi", expert_dims=expert_dims)
     if context.train and self.dropout_rate != 0.0:
       h = mtf.dropout(h, 1.0 - self.dropout_rate,
                       noise_shape=h.shape - context.length_dim)
     return mtf.layers.dense(h, io_channels, use_bias=False, activation=None,
                             variable_dtype=context.variable_dtype,
+                            reduced_dims=h.shape.dims[-1:],
                             name="wo", expert_dims=expert_dims)