mesh_tensorflow/transformer - fixes to incremental decoding.

PiperOrigin-RevId: 224011242

Author: nshazeer

mesh_tensorflow/transformer/moe.py

@@ -62,6 +62,13 @@ def __init__(self,
 
   def call(self, context, x, losses=None):
     """Call the layer."""
+    has_length_dim = context.length_dim in x.shape.dims
+    if not has_length_dim:
+      x_shape = x.shape
+      shape_with_length = mtf.Shape(
+          x_shape.dims[:-1] + [mtf.Dimension("length", 1)]
+          + x_shape.dims[-1:])
+      x = mtf.reshape(x, shape_with_length)
     y, loss = transformer_moe_layer_v1(
         x,
         context.model_dim,
@@ -70,6 +77,8 @@ def call(self, context, x, losses=None):
         context.variable_dtype)
     if context.losses is not None:
       context.losses.append(loss)
+    if not has_length_dim:
+      y = mtf.reshape(y, x_shape)
     return y
 
 
@@ -111,6 +120,13 @@ def __init__(self,
 
   def call(self, context, x, losses=None):
     """Call the layer."""
+    has_length_dim = context.length_dim in x.shape.dims
+    if not has_length_dim:
+      x_shape = x.shape
+      shape_with_length = mtf.Shape(
+          x_shape.dims[:-1] + [mtf.Dimension("length", 1)]
+          + x_shape.dims[-1:])
+      x = mtf.reshape(x, shape_with_length)
     y, loss = transformer_moe_layer_v2(
         x,
         context.model_dim,
@@ -119,6 +135,8 @@ def call(self, context, x, losses=None):
         context.variable_dtype)
     if context.losses is not None:
       context.losses.append(loss)
+    if not has_length_dim:
+      y = mtf.reshape(y, x_shape)
     return y
 
 

mesh_tensorflow/transformer/transformer_layers.py

@@ -153,9 +153,7 @@ def call(self, context, x, losses=None):
     if context.mode == "incremental":
       prev_k, prev_v = context.next_states(2)
       y, new_k, new_v = mtf.layers.masked_local_attention_1d_incremental(
-          x, prev_k, prev_v,
-          context.position, context.master_dtype, context.slice_dtype,
-          params=params)
+          x, prev_k, prev_v, context.position, params=params)
       context.new_states.extend([new_k, new_v])
       return y
     else: