Fix ConformerLayer's adapter_tpl for StreamStep

ConformerLayer.StreamStep doesn't respect p.adapter_tpl, while Fprop respects
it.

PiperOrigin-RevId: 491695042

Author: ds-hwang

lingvo/core/conformer_layer.py

@@ -1013,13 +1013,13 @@ def zero_state(self, batch_size):
     else:
       return py_utils.NestedMap()
 
-  def StreamStep(self, theta, inputs, paddings, state0):
+  def StreamStep(self, theta, in_nmap, state0):
     """Streams t steps.
 
     Args:
       theta: A NestedMap of read-only layer params.
-      inputs: A tensor of shape [b, t, d].
-      paddings: A 0/1 valued tensor of shape [b, t].
+      in_nmap: A NestedMap including features tensor of shape [b, t, d], and
+        paddings tensor of shape [b, t].
       state0: A NestedMap of tensors of the same struct as returned by
         zero_state().
 
@@ -1033,7 +1033,7 @@ def StreamStep(self, theta, inputs, paddings, state0):
     assert not p.remat
 
     with tf.name_scope(f'{p.name}/StreamStep'):
-      features, aux_loss = inputs, None
+      features, paddings, aux_loss = in_nmap.features, in_nmap.paddings, None
 
       if self.has_fflayer_start:
         features, paddings, aux_loss = self._MoeOrFFLayer(
@@ -1065,6 +1065,12 @@ def StreamStep(self, theta, inputs, paddings, state0):
                                                  paddings, aux_loss)
       outputs = self.final_ln.FProp(theta.final_ln, features)
 
+      if p.adapter_tpl:
+        adapter_in_map = in_nmap.DeepCopy()
+        adapter_in_map.features, adapter_in_map.padding = outputs, paddings
+        adapter_out_nmap = self.adapter.FProp(theta.adapter, adapter_in_map)
+        outputs, paddings = adapter_out_nmap.features, adapter_out_nmap.paddings
+
       state1 = py_utils.NestedMap(
           lconv_state=lconv_state1, atten_state=atten_state1)
       return outputs, paddings, state1

lingvo/core/conformer_layer_test.py

@@ -626,7 +626,8 @@ def _FProp(self, layer, inputs, paddings):
                        py_utils.NestedMap(features=inputs, paddings=paddings))
 
   def _StreamStep(self, layer, step_inputs, step_paddings, state):
-    return layer.StreamStep(layer.theta, step_inputs, step_paddings, state)
+    in_nmap = py_utils.NestedMap(features=step_inputs, paddings=step_paddings)
+    return layer.StreamStep(layer.theta, in_nmap, state)
 
   def _GetFPropOutput(self, fprop_out):
     return fprop_out.features, fprop_out.paddings