# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Encoders for the speech model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

from six.moves import range
from six.moves import zip
import tensorflow as tf

from tensorflow.python.ops import inplace_ops
from lingvo.core import base_encoder
from lingvo.core import base_layer
from lingvo.core import layers
from lingvo.core import plot
from lingvo.core import py_utils
from lingvo.core import rnn_cell
from lingvo.core import rnn_layers
from lingvo.core import summary_utils
from lingvo.core import model_helper

ConvLSTMBlock = collections.namedtuple('ConvLSTMBlock', ('rnn', 'cnn'))


class AsrEncoder(base_encoder.BaseEncoder):
  """Speech encoder version 1."""

  @classmethod
  def Params(cls):
    """Configs for AsrEncoder."""
    p = super(AsrEncoder, cls).Params()
    p.Define('lstm_tpl', rnn_cell.LSTMCellSimple.Params(),
             'Configs template for the RNN layer.')
    p.Define('cnn_tpl', layers.ConvLayer.Params(),
             'Configs template for the conv layer.')
    p.Define('proj_tpl', layers.ProjectionLayer.Params(),
             'Configs template for the projection layer.')
    p.Define(
        'highway_skip', False,
        'If set, residual connections from different layers are gated. '
        'Will only be used if residual_start is enabled.')
    p.Define('highway_skip_tpl', layers.HighwaySkipLayer.Params(),
             'Configs template for the highway skip layer.')
    p.Define('conv_lstm_tpl', rnn_cell.ConvLSTMCell.Params(),
             'Configs template for ConvLSTMCell.')
    p.Define(
        'after_conv_lstm_cnn_tpl', layers.ConvLayer.Params(),
        'Configs template for the cnn layer immediately follow the'
        ' convlstm layer.')
    p.Define('conv_filter_shapes', None, 'Filter shapes for each conv layer.')
    p.Define('conv_filter_strides', None, 'Filter strides for each conv layer.')
    p.Define('input_shape', [None, None, None, None],
             'Shape of the input. This should a TensorShape with rank 4.')
    p.Define('lstm_cell_size', 256, 'LSTM cell size for the RNN layer.')
    p.Define('num_cnn_layers', 2, 'Number of conv layers to create.')
    p.Define('num_conv_lstm_layers', 1, 'Number of conv lstm layers to create.')
    p.Define('num_lstm_layers', 3, 'Number of rnn layers to create')
    p.Define('project_lstm_output', True,
             'Include projection layer after each encoder LSTM layer.')
    p.Define('pad_steps', 6,
             'Extra zero-padded timesteps to add to the input sequence. ')
    p.Define(
        'residual_start', 0, 'Start residual connections from this lstm layer. '
        'Disabled if 0 or greater than num_lstm_layers.')
    p.Define('residual_stride', 1,
             'Number of lstm layers to skip per residual connection.')
    p.Define(
        'bidi_rnn_type', 'func', 'Options: func, native_cudnn. '
        'func: BidirectionalFRNN, '
        'native_cudnn: BidirectionalNativeCuDNNLSTM.')

    # TODO(yonghui): Maybe move those configs to a separate file.
    # Set some reasonable default values.
    #
    # NOTE(yonghui): The default config below assumes the following encoder
    # architecture:
    #
    #   cnn/batch-norm/relu ->
    #   cnn/batch-norm/relu ->
    #   bidirectional conv-lstm ->
    #   cnn/batch-norm/relu
    #   bidirectional lstm ->
    #   projection/batch-norm/relu ->
    #   bidirectional lstm ->
    #   projection/batch-norm/relu ->
    #   bidirectional lstm
    #
    # Default config for the rnn layer.
    p.lstm_tpl.params_init = py_utils.WeightInit.Uniform(0.1)

    # Default config for the convolution layer.
    p.input_shape = [None, None, 80, 3]
    p.conv_filter_shapes = [(3, 3, 3, 32), (3, 3, 32, 32)]
    p.conv_filter_strides = [(2, 2), (2, 2)]
    p.cnn_tpl.params_init = py_utils.WeightInit.TruncatedGaussian(0.1)
    # TODO(yonghui): Disable variational noise logic.
    # NOTE(yonghui): Fortunately, variational noise logic is currently not
    # implemented for ConvLayer yet (as of sep 22, 2016).

    # Default config for the projection layer.
    p.proj_tpl.params_init = py_utils.WeightInit.TruncatedGaussian(0.1)
    # TODO(yonghui): Disable variational noise logic.
    # NOTE(yonghui): Fortunately, variational noise logic is currently not
    # implemented for ProjectionLayer yet (as of sep 22, 2016).

    p.conv_lstm_tpl.filter_shape = [1, 3]  # height (time), width (frequency)
    p.conv_lstm_tpl.inputs_shape = [None, None, None, None]
    p.conv_lstm_tpl.cell_shape = [None, None, None, None]
    p.conv_lstm_tpl.params_init = py_utils.WeightInit.TruncatedGaussian(0.1)
    p.after_conv_lstm_cnn_tpl.filter_shape = [3, 3, None, None]
    p.after_conv_lstm_cnn_tpl.params_init = (
        py_utils.WeightInit.TruncatedGaussian(0.1))
    p.after_conv_lstm_cnn_tpl.filter_stride = [1, 1]
    return p

  @base_layer.initializer
  def __init__(self, params):
    super(AsrEncoder, self).__init__(params)
    p = self.params
    assert p.packed_input is False, ('Packed inputs are not yet supported for '
                                     'AsrEncoder.')
    name = p.name

    with tf.variable_scope(name):
      # First create the conv layers.

      assert p.num_cnn_layers == len(p.conv_filter_shapes)
      assert p.num_cnn_layers == len(p.conv_filter_strides)
      params_conv_layers = []
      for i in range(p.num_cnn_layers):
        conv_p = p.cnn_tpl.Copy()
        conv_p.name = 'conv_L%d' % (i)
        conv_p.filter_shape = p.conv_filter_shapes[i]
        conv_p.filter_stride = p.conv_filter_strides[i]
        conv_p.is_eval = p.is_eval
        params_conv_layers.append(conv_p)
      self.CreateChildren('conv', params_conv_layers)

      conv_output_shape = tf.TensorShape(p.input_shape)
      for i in range(p.num_cnn_layers):
        conv_output_shape = self.conv[i].OutShape(conv_output_shape)
      conv_output_shape = conv_output_shape.as_list()
      assert len(conv_output_shape) == 4  # batch, height, width, channel.

      params_conv_lstm_rnn = []
      params_conv_lstm_cnn = []
      for i in range(p.num_conv_lstm_layers):
        # NOTE(yonghui): We assume that output from ConvLSTMBlock has the same
        # shape as its input.
        _, _, width, in_channel = conv_output_shape
        f_conv_lstm_p = p.conv_lstm_tpl.Copy()
        f_conv_lstm_p.name = 'f_conv_lstm_%d' % (i)
        f_conv_lstm_p.inputs_shape = [None, 1, width, in_channel]
        f_conv_lstm_p.cell_shape = [None, 1, width, in_channel]
        b_conv_lstm_p = f_conv_lstm_p.Copy()
        b_conv_lstm_p.name = 'b_conv_lstm_%d' % (i)
        conv_lstm_rnn_p = self.CreateConvLstmLayerParams()
        conv_lstm_rnn_p.name = 'conv_lstm_rnn'
        conv_lstm_rnn_p.fwd = f_conv_lstm_p
        conv_lstm_rnn_p.bak = b_conv_lstm_p
        params_conv_lstm_rnn.append(conv_lstm_rnn_p)
        cnn_p = p.after_conv_lstm_cnn_tpl.Copy()
        cnn_p.name = 'conv_lstm_cnn_%d' % (i)
        cnn_p.filter_shape[2] = 2 * in_channel
        cnn_p.filter_shape[3] = in_channel
        params_conv_lstm_cnn.append(cnn_p)
        # TODO(yonghui): Refactor ConvLSTMBlock into a layer.
      self.CreateChildren('conv_lstm_rnn', params_conv_lstm_rnn)
      self.CreateChildren('conv_lstm_cnn', params_conv_lstm_cnn)

      (self._first_lstm_input_dim,
       self._first_lstm_input_dim_pad) = self.FirstLstmLayerInputDimAndPadding(
           conv_output_shape, pad_to_multiple=16)

      # Now create all the rnn layers and projection layers.
      # TODO(yonghui): take care of device placement.
      params_rnn_layers = []
      params_proj_layers = []
      params_highway_skip_layers = []
      for i in range(p.num_lstm_layers):
        if i == 0:
          input_dim = self._first_lstm_input_dim
        else:
          input_dim = 2 * p.lstm_cell_size
        forward_p = p.lstm_tpl.Copy()
        forward_p.name = 'fwd_rnn_L%d' % (i)
        forward_p.num_input_nodes = input_dim
        forward_p.num_output_nodes = p.lstm_cell_size
        backward_p = forward_p.Copy()
        backward_p.name = 'bak_rnn_L%d' % (i)
        rnn_p = self.CreateBidirectionalRNNParams(forward_p, backward_p)
        rnn_p.name = 'brnn_L%d' % (i)
        params_rnn_layers.append(rnn_p)

        if p.project_lstm_output and (i < p.num_lstm_layers - 1):
          proj_p = p.proj_tpl.Copy()
          proj_p.input_dim = 2 * p.lstm_cell_size
          proj_p.output_dim = 2 * p.lstm_cell_size
          proj_p.name = 'proj_L%d' % (i)
          proj_p.is_eval = p.is_eval
          params_proj_layers.append(proj_p)

        # add the skip layers
        residual_index = i - p.residual_start + 1
        if p.residual_start > 0 and residual_index >= 0 and p.highway_skip:
          highway_skip = p.highway_skip_tpl.Copy()
          highway_skip.name = 'enc_hwskip_%d' % len(params_highway_skip_layers)
          highway_skip.input_dim = 2 * p.lstm_cell_size
          params_highway_skip_layers.append(highway_skip)
      self.CreateChildren('rnn', params_rnn_layers)
      self.CreateChildren('proj', params_proj_layers)
      self.CreateChildren('highway_skip', params_highway_skip_layers)

  @property
  def _use_functional(self):
    return True

  def CreateBidirectionalRNNParams(self, forward_p, backward_p):
    return model_helper.CreateBidirectionalRNNParams(self.params, forward_p,
                                                     backward_p)

  def CreateConvLstmLayerParams(self):
    return rnn_layers.BidirectionalFRNN.Params()

  def FirstLstmLayerInputDimAndPadding(self,
                                       conv_output_shape,
                                       pad_to_multiple=16):
    lstm_input_shape = conv_output_shape
    # Makes sure the lstm input dims is multiple of 16 (alignment
    # requirement from FRNN).
    first_lstm_input_dim_unpadded = lstm_input_shape[2] * lstm_input_shape[3]

    if self._use_functional and (first_lstm_input_dim_unpadded % pad_to_multiple
                                 != 0):
      first_lstm_input_dim = int(
          (first_lstm_input_dim_unpadded + pad_to_multiple - 1) /
          pad_to_multiple) * pad_to_multiple
    else:
      first_lstm_input_dim = first_lstm_input_dim_unpadded

    first_lstm_input_dim_padding = (
        first_lstm_input_dim - first_lstm_input_dim_unpadded)
    return first_lstm_input_dim, first_lstm_input_dim_padding

  @property
  def supports_streaming(self):
    return False

  def zero_state(self, batch_size):
    return py_utils.NestedMap()

  def FProp(self, theta, batch, state0=None):
    """Encodes source as represented by 'inputs' and 'paddings'.

    Args:
      theta: A NestedMap object containing weights' values of this
        layer and its children layers.
      batch: A NestedMap with fields:
        src_inputs - The inputs tensor. It is expected to be of shape [batch,
            time, feature_dim, channels].
        paddings - The paddings tensor. It is expected to be of shape [batch,
          time].
      state0: Recurrent input state. Not supported/ignored by this encoder.

    Returns:
      (outputs, out_paddings, state1) tuple. Outputs is of the shape
      [time, batch, depth], and out_paddings is of the shape [time, batch]
    """
    p = self.params
    inputs, paddings = batch.src_inputs, batch.paddings
    with tf.name_scope(p.name):
      # Add a few extra padded timesteps at the end. This is for ensuring the
      # correctness of the conv-layers at the edges.
      if p.pad_steps > 0:
        # inplace_update() is not supported by TPU for now. Since we have done
        # padding on the input_generator, we may avoid this additional padding.
        assert not py_utils.use_tpu()
        inputs_pad = tf.zeros(
            inplace_ops.inplace_update(tf.shape(inputs), 1, p.pad_steps),
            inputs.dtype)
        paddings_pad = tf.ones(
            inplace_ops.inplace_update(tf.shape(paddings), 1, p.pad_steps),
            paddings.dtype)
        inputs = tf.concat([inputs, inputs_pad], 1, name='inputs')
        paddings = tf.concat([paddings, paddings_pad], 1)

      def ReshapeForPlot(tensor, padding, name):
        """Transposes and flattens channels to [batch, dim, seq_len] shape."""
        # Flatten any dimensions beyond the third into the third.
        batch_size = tf.shape(tensor)[0]
        max_len = tf.shape(tensor)[1]
        plot_tensor = tf.reshape(tensor, [batch_size, max_len, -1])
        plot_tensor = tf.transpose(plot_tensor, [0, 2, 1], name=name)
        return (plot_tensor, summary_utils.SequenceLength(padding))

      plots = [
          ReshapeForPlot(
              tf.transpose(inputs, [0, 1, 3, 2]), paddings, 'inputs')
      ]

      conv_out = inputs
      out_padding = paddings
      for i, conv_layer in enumerate(self.conv):
        conv_out, out_padding = conv_layer.FProp(theta.conv[i], conv_out,
                                                 out_padding)
        plots.append(
            ReshapeForPlot(
                tf.transpose(conv_out, [0, 1, 3, 2]), out_padding,
                'conv_%d_out' % i))

      def TransposeFirstTwoDims(t):
        first_dim = tf.shape(t)[0]
        second_dim = tf.shape(t)[1]
        t_new = tf.transpose(
            tf.reshape(t, [first_dim, second_dim, -1]), [1, 0, 2])
        t_shape_new = tf.concat([[second_dim], [first_dim], tf.shape(t)[2:]], 0)
        return tf.reshape(t_new, t_shape_new)

      # Now the conv-lstm part.
      conv_lstm_out = conv_out
      conv_lstm_out_padding = out_padding
      for i, (rnn, cnn) in enumerate(
          zip(self.conv_lstm_rnn, self.conv_lstm_cnn)):
        conv_lstm_in = conv_lstm_out
        # Move time dimension to be the first.
        conv_lstm_in = TransposeFirstTwoDims(conv_lstm_in)
        conv_lstm_in = tf.expand_dims(conv_lstm_in, 2)
        conv_lstm_in_padding = tf.expand_dims(
            tf.transpose(conv_lstm_out_padding), 2)
        lstm_out = rnn.FProp(theta.conv_lstm_rnn[i], conv_lstm_in,
                             conv_lstm_in_padding)
        # Move time dimension to be the second.
        cnn_in = TransposeFirstTwoDims(lstm_out)
        cnn_in = tf.squeeze(cnn_in, 2)
        cnn_in_padding = conv_lstm_out_padding
        cnn_out, cnn_out_padding = cnn.FProp(theta.conv_lstm_cnn[i], cnn_in,
                                             cnn_in_padding)
        conv_lstm_out, conv_lstm_out_padding = cnn_out, cnn_out_padding
        plots.append(
            ReshapeForPlot(conv_lstm_out, conv_lstm_out_padding,
                           'conv_lstm_%d_out' % i))

      # Need to do a reshape before starting the rnn layers.
      conv_lstm_out = py_utils.HasRank(conv_lstm_out, 4)
      conv_lstm_out_shape = tf.shape(conv_lstm_out)
      new_shape = tf.concat([conv_lstm_out_shape[:2], [-1]], 0)
      conv_lstm_out = tf.reshape(conv_lstm_out, new_shape)
      if self._first_lstm_input_dim_pad:
        conv_lstm_out = tf.pad(
            conv_lstm_out,
            [[0, 0], [0, 0], [0, self._first_lstm_input_dim_pad]])

      conv_lstm_out = py_utils.HasShape(conv_lstm_out,
                                        [-1, -1, self._first_lstm_input_dim])

      # Transpose to move the time dimension to be the first.
      rnn_in = tf.transpose(conv_lstm_out, [1, 0, 2])
      rnn_padding = tf.expand_dims(tf.transpose(conv_lstm_out_padding), 2)
      # rnn_in is of shape [time, batch, depth]
      # rnn_padding is of shape [time, batch, 1]

      # Now the rnn layers.
      num_skips = 0
      for i in range(p.num_lstm_layers):
        rnn_out = self.rnn[i].FProp(theta.rnn[i], rnn_in, rnn_padding)
        residual_index = i - p.residual_start + 1
        if p.residual_start > 0 and residual_index >= 0:
          if residual_index % p.residual_stride == 0:
            residual_in = rnn_in
          if residual_index % p.residual_stride == p.residual_stride - 1:
            # Highway skip connection.
            if p.highway_skip:
              rnn_out = self.highway_skip[num_skips].FProp(
                  theta.highway_skip[num_skips], residual_in, rnn_out)
              num_skips += 1
            else:
              # Residual skip connection.
              rnn_out += py_utils.HasShape(residual_in, tf.shape(rnn_out))
        if p.project_lstm_output and (i < p.num_lstm_layers - 1):
          # Projection layers.
          rnn_out = self.proj[i].FProp(theta.proj[i], rnn_out, rnn_padding)
        if i == p.num_lstm_layers - 1:
          rnn_out *= (1.0 - rnn_padding)
        plots.append(
            ReshapeForPlot(
                tf.transpose(rnn_out, [1, 0, 2]),
                tf.transpose(rnn_padding, [1, 0, 2]), 'rnn_%d_out' % i))
        rnn_in = rnn_out
      final_out = rnn_in

      if self.cluster.add_summary:
        fig = plot.MatplotlibFigureSummary(
            'encoder_example', figsize=(8, len(plots) * 3.5))

        # Order layers from bottom to top.
        plots.reverse()
        for tensor, seq_len in plots:
          fig.AddSubplot(
              [tensor, seq_len],
              summary_utils.TrimPaddingAndPlotSequence,
              title=tensor.name,
              xlabel='Time')
        fig.Finalize()

      rnn_padding = tf.squeeze(rnn_padding, [2])
      return final_out, rnn_padding, py_utils.NestedMap()

[((36, 16, 36, 71), 'collections.namedtuple', 'collections.namedtuple', ({(36, 39, 36, 54): '"""ConvLSTMBlock"""', (36, 56, 36, 70): "('rnn', 'cnn')"}, {}), "('ConvLSTMBlock', ('rnn', 'cnn'))", False, 'import collections\n'), ((103, 29, 103, 61), 'lingvo.core.py_utils.WeightInit.Uniform', 'py_utils.WeightInit.Uniform', ({(103, 57, 103, 60): '0.1'}, {}), '(0.1)', False, 'from lingvo.core import py_utils\n'), ((109, 28, 109, 70), 'lingvo.core.py_utils.WeightInit.TruncatedGaussian', 'py_utils.WeightInit.TruncatedGaussian', ({(109, 66, 109, 69): '0.1'}, {}), '(0.1)', False, 'from lingvo.core import py_utils\n'), ((115, 29, 115, 71), 'lingvo.core.py_utils.WeightInit.TruncatedGaussian', 'py_utils.WeightInit.TruncatedGaussian', ({(115, 67, 115, 70): '0.1'}, {}), '(0.1)', False, 'from lingvo.core import py_utils\n'), ((123, 34, 123, 76), 'lingvo.core.py_utils.WeightInit.TruncatedGaussian', 'py_utils.WeightInit.TruncatedGaussian', ({(123, 72, 123, 75): '0.1'}, {}), '(0.1)', False, 'from lingvo.core import py_utils\n'), ((126, 8, 126, 50), 'lingvo.core.py_utils.WeightInit.TruncatedGaussian', 'py_utils.WeightInit.TruncatedGaussian', ({(126, 46, 126, 49): '0.1'}, {}), '(0.1)', False, 'from lingvo.core import py_utils\n'), ((233, 11, 234, 64), 'lingvo.core.model_helper.CreateBidirectionalRNNParams', 'model_helper.CreateBidirectionalRNNParams', ({(233, 53, 233, 64): 'self.params', (233, 66, 233, 75): 'forward_p', (234, 53, 234, 63): 'backward_p'}, {}), '(self.params, forward_p, backward_p)', False, 'from lingvo.core import model_helper\n'), ((237, 11, 237, 48), 'lingvo.core.rnn_layers.BidirectionalFRNN.Params', 'rnn_layers.BidirectionalFRNN.Params', ({}, {}), '()', False, 'from lingvo.core import rnn_layers\n'), ((264, 11, 264, 31), 'lingvo.core.py_utils.NestedMap', 'py_utils.NestedMap', ({}, {}), '()', False, 'from lingvo.core import py_utils\n'), ((46, 25, 46, 57), 'lingvo.core.rnn_cell.LSTMCellSimple.Params', 'rnn_cell.LSTMCellSimple.Params', ({}, {}), '()', False, 'from lingvo.core import rnn_cell\n'), ((48, 24, 48, 49), 'lingvo.core.layers.ConvLayer.Params', 'layers.ConvLayer.Params', ({}, {}), '()', False, 'from lingvo.core import layers\n'), ((50, 25, 50, 56), 'lingvo.core.layers.ProjectionLayer.Params', 'layers.ProjectionLayer.Params', ({}, {}), '()', False, 'from lingvo.core import layers\n'), ((56, 33, 56, 65), 'lingvo.core.layers.HighwaySkipLayer.Params', 'layers.HighwaySkipLayer.Params', ({}, {}), '()', False, 'from lingvo.core import layers\n'), ((58, 30, 58, 60), 'lingvo.core.rnn_cell.ConvLSTMCell.Params', 'rnn_cell.ConvLSTMCell.Params', ({}, {}), '()', False, 'from lingvo.core import rnn_cell\n'), ((61, 35, 61, 60), 'lingvo.core.layers.ConvLayer.Params', 'layers.ConvLayer.Params', ({}, {}), '()', False, 'from lingvo.core import layers\n'), ((138, 9, 138, 32), 'tensorflow.variable_scope', 'tf.variable_scope', ({(138, 27, 138, 31): 'name'}, {}), '(name)', True, 'import tensorflow as tf\n'), ((144, 15, 144, 38), 'six.moves.range', 'range', ({(144, 21, 144, 37): 'p.num_cnn_layers'}, {}), '(p.num_cnn_layers)', False, 'from six.moves import range\n'), ((153, 26, 153, 55), 'tensorflow.TensorShape', 'tf.TensorShape', ({(153, 41, 153, 54): 'p.input_shape'}, {}), '(p.input_shape)', True, 'import tensorflow as tf\n'), ((154, 15, 154, 38), 'six.moves.range', 'range', ({(154, 21, 154, 37): 'p.num_cnn_layers'}, {}), '(p.num_cnn_layers)', False, 'from six.moves import range\n'), ((161, 15, 161, 44), 'six.moves.range', 'range', ({(161, 21, 161, 43): 'p.num_conv_lstm_layers'}, {}), '(p.num_conv_lstm_layers)', False, 'from six.moves import range\n'), ((194, 15, 194, 39), 'six.moves.range', 'range', ({(194, 21, 194, 38): 'p.num_lstm_layers'}, {}), '(p.num_lstm_layers)', False, 'from six.moves import range\n'), ((285, 9, 285, 30), 'tensorflow.name_scope', 'tf.name_scope', ({(285, 23, 285, 29): 'p.name'}, {}), '(p.name)', True, 'import tensorflow as tf\n'), ((358, 22, 358, 56), 'lingvo.core.py_utils.HasRank', 'py_utils.HasRank', ({(358, 39, 358, 52): 'conv_lstm_out', (358, 54, 358, 55): '4'}, {}), '(conv_lstm_out, 4)', False, 'from lingvo.core import py_utils\n'), ((359, 28, 359, 51), 'tensorflow.shape', 'tf.shape', ({(359, 37, 359, 50): 'conv_lstm_out'}, {}), '(conv_lstm_out)', True, 'import tensorflow as tf\n'), ((360, 18, 360, 63), 'tensorflow.concat', 'tf.concat', ({(360, 28, 360, 59): '[conv_lstm_out_shape[:2], [-1]]', (360, 61, 360, 62): '0'}, {}), '([conv_lstm_out_shape[:2], [-1]], 0)', True, 'import tensorflow as tf\n'), ((361, 22, 361, 58), 'tensorflow.reshape', 'tf.reshape', ({(361, 33, 361, 46): 'conv_lstm_out', (361, 48, 361, 57): 'new_shape'}, {}), '(conv_lstm_out, new_shape)', True, 'import tensorflow as tf\n'), ((367, 22, 368, 77), 'lingvo.core.py_utils.HasShape', 'py_utils.HasShape', ({(367, 40, 367, 53): 'conv_lstm_out', (368, 40, 368, 76): '[-1, -1, self._first_lstm_input_dim]'}, {}), '(conv_lstm_out, [-1, -1, self._first_lstm_input_dim])', False, 'from lingvo.core import py_utils\n'), ((371, 15, 371, 53), 'tensorflow.transpose', 'tf.transpose', ({(371, 28, 371, 41): 'conv_lstm_out', (371, 43, 371, 52): '[1, 0, 2]'}, {}), '(conv_lstm_out, [1, 0, 2])', True, 'import tensorflow as tf\n'), ((378, 15, 378, 39), 'six.moves.range', 'range', ({(378, 21, 378, 38): 'p.num_lstm_layers'}, {}), '(p.num_lstm_layers)', False, 'from six.moves import range\n'), ((419, 20, 419, 48), 'tensorflow.squeeze', 'tf.squeeze', ({(419, 31, 419, 42): 'rnn_padding', (419, 44, 419, 47): '[2]'}, {}), '(rnn_padding, [2])', True, 'import tensorflow as tf\n'), ((298, 17, 298, 66), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((299, 19, 299, 57), 'tensorflow.concat', 'tf.concat', ({(299, 29, 299, 53): '[paddings, paddings_pad]', (299, 55, 299, 56): '1'}, {}), '([paddings, paddings_pad], 1)', True, 'import tensorflow as tf\n'), ((306, 22, 306, 67), 'tensorflow.reshape', 'tf.reshape', ({(306, 33, 306, 39): 'tensor', (306, 41, 306, 66): '[batch_size, max_len, -1]'}, {}), '(tensor, [batch_size, max_len, -1])', True, 'import tensorflow as tf\n'), ((307, 22, 307, 69), 'tensorflow.transpose', 'tf.transpose', (), '', True, 'import tensorflow as tf\n'), ((331, 15, 331, 45), 'tensorflow.reshape', 'tf.reshape', ({(331, 26, 331, 31): 't_new', (331, 33, 331, 44): 't_shape_new'}, {}), '(t_new, t_shape_new)', True, 'import tensorflow as tf\n'), ((337, 10, 337, 53), 'six.moves.zip', 'zip', ({(337, 14, 337, 32): 'self.conv_lstm_rnn', (337, 34, 337, 52): 'self.conv_lstm_cnn'}, {}), '(self.conv_lstm_rnn, self.conv_lstm_cnn)', False, 'from six.moves import zip\n'), ((341, 23, 341, 54), 'tensorflow.expand_dims', 'tf.expand_dims', ({(341, 38, 341, 50): 'conv_lstm_in', (341, 52, 341, 53): '2'}, {}), '(conv_lstm_in, 2)', True, 'import tensorflow as tf\n'), ((348, 17, 348, 38), 'tensorflow.squeeze', 'tf.squeeze', ({(348, 28, 348, 34): 'cnn_in', (348, 36, 348, 37): '2'}, {}), '(cnn_in, 2)', True, 'import tensorflow as tf\n'), ((363, 24, 365, 66), 'tensorflow.pad', 'tf.pad', ({(364, 12, 364, 25): 'conv_lstm_out', (365, 12, 365, 65): '[[0, 0], [0, 0], [0, self._first_lstm_input_dim_pad]]'}, {}), '(conv_lstm_out, [[0, 0], [0, 0], [0, self._first_lstm_input_dim_pad]])', True, 'import tensorflow as tf\n'), ((372, 35, 372, 70), 'tensorflow.transpose', 'tf.transpose', ({(372, 48, 372, 69): 'conv_lstm_out_padding'}, {}), '(conv_lstm_out_padding)', True, 'import tensorflow as tf\n'), ((420, 37, 420, 57), 'lingvo.core.py_utils.NestedMap', 'py_utils.NestedMap', ({}, {}), '()', False, 'from lingvo.core import py_utils\n'), ((291, 19, 291, 37), 'lingvo.core.py_utils.use_tpu', 'py_utils.use_tpu', ({}, {}), '()', False, 'from lingvo.core import py_utils\n'), ((304, 21, 304, 37), 'tensorflow.shape', 'tf.shape', ({(304, 30, 304, 36): 'tensor'}, {}), '(tensor)', True, 'import tensorflow as tf\n'), ((305, 18, 305, 34), 'tensorflow.shape', 'tf.shape', ({(305, 27, 305, 33): 'tensor'}, {}), '(tensor)', True, 'import tensorflow as tf\n'), ((308, 29, 308, 66), 'lingvo.core.summary_utils.SequenceLength', 'summary_utils.SequenceLength', ({(308, 58, 308, 65): 'padding'}, {}), '(padding)', False, 'from lingvo.core import summary_utils\n'), ((312, 14, 312, 48), 'tensorflow.transpose', 'tf.transpose', ({(312, 27, 312, 33): 'inputs', (312, 35, 312, 47): '[0, 1, 3, 2]'}, {}), '(inputs, [0, 1, 3, 2])', True, 'import tensorflow as tf\n'), ((326, 20, 326, 31), 'tensorflow.shape', 'tf.shape', ({(326, 29, 326, 30): 't'}, {}), '(t)', True, 'import tensorflow as tf\n'), ((327, 21, 327, 32), 'tensorflow.shape', 'tf.shape', ({(327, 30, 327, 31): 't'}, {}), '(t)', True, 'import tensorflow as tf\n'), ((329, 12, 329, 54), 'tensorflow.reshape', 'tf.reshape', ({(329, 23, 329, 24): 't', (329, 26, 329, 53): '[first_dim, second_dim, -1]'}, {}), '(t, [first_dim, second_dim, -1])', True, 'import tensorflow as tf\n'), ((343, 12, 343, 47), 'tensorflow.transpose', 'tf.transpose', ({(343, 25, 343, 46): 'conv_lstm_out_padding'}, {}), '(conv_lstm_out_padding)', True, 'import tensorflow as tf\n'), ((293, 39, 293, 55), 'tensorflow.shape', 'tf.shape', ({(293, 48, 293, 54): 'inputs'}, {}), '(inputs)', True, 'import tensorflow as tf\n'), ((296, 39, 296, 57), 'tensorflow.shape', 'tf.shape', ({(296, 48, 296, 56): 'paddings'}, {}), '(paddings)', True, 'import tensorflow as tf\n'), ((322, 16, 322, 52), 'tensorflow.transpose', 'tf.transpose', ({(322, 29, 322, 37): 'conv_out', (322, 39, 322, 51): '[0, 1, 3, 2]'}, {}), '(conv_out, [0, 1, 3, 2])', True, 'import tensorflow as tf\n'), ((400, 16, 400, 48), 'tensorflow.transpose', 'tf.transpose', ({(400, 29, 400, 36): 'rnn_out', (400, 38, 400, 47): '[1, 0, 2]'}, {}), '(rnn_out, [1, 0, 2])', True, 'import tensorflow as tf\n'), ((401, 16, 401, 52), 'tensorflow.transpose', 'tf.transpose', ({(401, 29, 401, 40): 'rnn_padding', (401, 42, 401, 51): '[1, 0, 2]'}, {}), '(rnn_padding, [1, 0, 2])', True, 'import tensorflow as tf\n'), ((330, 60, 330, 71), 'tensorflow.shape', 'tf.shape', ({(330, 69, 330, 70): 't'}, {}), '(t)', True, 'import tensorflow as tf\n'), ((392, 56, 392, 73), 'tensorflow.shape', 'tf.shape', ({(392, 65, 392, 72): 'rnn_out'}, {}), '(rnn_out)', True, 'import tensorflow as tf\n')]