Use tf.slice instead of tf.strided_slice (via tf.Tensor.__getitem__ / indexing syntax) for state tensors.

This ensures that tensors in the output state have the same static shape as tensors in the input state. Before this change the equivalence is only guaranteed for the dynamic shape.

PiperOrigin-RevId: 480168486

Author: lingvo-bot

lingvo/core/batch_major_attention.py

@@ -2271,7 +2271,8 @@ def StreamStepAddSkipConnection(self, input_to_add, output, state0, state1):
     concat_input_to_add = tf.concat([state0.skip_conn_input, input_to_add],
                                     axis=1)
     final_output = output + concat_input_to_add[:, :seqlen]
-    state1.skip_conn_input = concat_input_to_add[:, seqlen:]
+    state1.skip_conn_input = tf.slice(concat_input_to_add, [0, seqlen, 0],
+                                      tf.shape(state0.skip_conn_input))
     return final_output, state1
 
   def _StreamStepDimensions(self, inputs):
@@ -2664,12 +2665,14 @@ def _StreamStepDynamicLength(self, theta, query_vec, query_paddings, key_vec,
       output = self.post.FProp(theta.post, output)
 
       state1 = py_utils.NestedMap(
-          key=key[:, k:, :, :],
-          value=value[:, k:, :, :],
-          masks=state_masks[:, k:])
+          key=tf.slice(key, [0, k, 0, 0], tf.shape(state0.key)),
+          value=tf.slice(value, [0, k, 0, 0], tf.shape(state0.value)),
+          masks=tf.slice(state_masks, [0, k], tf.shape(state0.masks)))
       if p.right_context > 0:
-        state1.query = concat_query[:, q:]
-        state1.out_masks = concat_out_masks[:, q:]
+        state1.query = tf.slice(concat_query, [0, q, 0, 0],
+                                tf.shape(state0.query))
+        state1.out_masks = tf.slice(concat_out_masks, [0, q],
+                                    tf.shape(state0.out_masks))
       return output, out_paddings, state1
 
   @classmethod

lingvo/core/conv_layers_with_time_padding.py

@@ -596,7 +596,8 @@ def StreamStep(self, theta, inputs, paddings, state0):
           padding='VALID')
       if p.bias:
         outputs = tf.nn.bias_add(outputs, theta.b)
-      new_context = concat_inputs[:, q:]
+      new_context = tf.slice(concat_inputs, [0, q, 0, 0],
+                             tf.shape(state0.context))
       return outputs, paddings, py_utils.NestedMap(context=new_context)
 
 
@@ -773,7 +774,8 @@ def StreamStep(self, theta, inputs, paddings, state0):
           padding='VALID')
       if p.bias:
         outputs = tf.nn.bias_add(outputs, theta.b)
-      new_context = concat_inputs[:, q:]
+      new_context = tf.slice(concat_inputs, [0, q, 0, 0],
+                             tf.shape(state0.context))
       return outputs, paddings, py_utils.NestedMap(context=new_context)