Python tensorflow.TensorArray() Examples

def tas_for_tensors(tensors, length):
  """Unstacks a set of Tensors into TensorArrays.

  Args:
    tensors: A potentially nested tuple or list of Tensors with length in the
      first dimension greater than or equal to the 'length' input argument.
    length: The desired length of the TensorArrays.
  Returns:
    tensorarrays: A potentially nested tuple or list of TensorArrays with the
      same structure as 'tensors'. Contains the result of unstacking each Tensor
      in 'tensors'.
  """
  def map_fn(x):
    ta = tf.TensorArray(x.dtype, length, name=x.name.split(':')[0] + '_ta')
    return ta.unstack(x[:length, :])
  return map_nested(map_fn, tensors) 

def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
    """Raises an exception if dimensions are known statically and can not be
    reshaped to [batch_size, beam_size, -1]."""
    reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
    assert len(shape.dims) > 0
    if batch_size is None or shape[0] is None:
        return True  # not statically known => no check
    if shape[0] == batch_size * beam_width:
        return True  # flattened, matching
    has_second_dim = shape.ndims >= 2 and shape[1] is not None
    if has_second_dim and shape[0] == batch_size and shape[1] == beam_width:
        return True  # non-flattened, matching
    # Otherwise we could not find a match and warn:
    tf.get_logger().warn(
        "TensorArray reordering expects elements to be "
        "reshapable to %s which is incompatible with the "
        "current shape %s. Consider setting "
        "reorder_tensor_arrays to False to disable TensorArray "
        "reordering during the beam search." % (reshaped_shape, shape)
    )
    return False 

def _maybe_merge_batch_beams(self, t, s):
        """Splits the tensor from a batch by beams into a batch of beams.

        More exactly, `t` is a tensor of dimension
        `[batch_size * beam_width] + s`, then we reshape it to
        `[batch_size, beam_width] + s`.

        Args:
          t: `Tensor` of dimension `[batch_size * beam_width] + s`.
          s: `Tensor`, Python int, or `TensorShape`.

        Returns:
          A reshaped version of t with shape `[batch_size, beam_width] + s`.

        Raises:
          ValueError:  If the rank of `t` is not statically known.
        """
        if isinstance(t, tf.TensorArray):
            return t
        _check_ndims(t)
        if t.shape.ndims >= 2:
            return self._merge_batch_beams(t, s)
        else:
            return t 

def mix_target_sequence(self, gold_token, predicted_token, training, top_k=5):
        """ to mix gold token and prediction
        param gold_token: true labels
        param predicted_token: predictions by first pass
        return: mix of the gold_token and predicted_token
        """
        mix_result = tf.TensorArray(
            tf.float32, size=1, dynamic_size=True, clear_after_read=False
        )
        for i in tf.range(tf.shape(gold_token)[-1]):
            if self.random_num([1]) > self.hparams.schedual_sampling_rate:# do schedual sampling
                selected_input = predicted_token[:, i, :]
                selected_idx = tf.nn.top_k(selected_input, top_k).indices
                embedding_input = self.y_net.layers[1](selected_idx, training=training)
                embedding_input = tf.reduce_mean(embedding_input, axis=1)
                mix_result = mix_result.write(i, embedding_input)
            else:
                selected_input = tf.reshape(gold_token[:, i], [-1, 1])
                embedding_input = self.y_net.layers[1](selected_input, training=training)
                mix_result = mix_result.write(i, embedding_input[:, 0, :])
        final_input = self.y_net.layers[2](tf.transpose(mix_result.stack(), [1, 0, 2]),
                                           training=training)
        final_input = self.y_net.layers[3](final_input, training=training)
        return final_input 

def tas_for_tensors(tensors, length):
  """Unstacks a set of Tensors into TensorArrays.

  Args:
    tensors: A potentially nested tuple or list of Tensors with length in the
      first dimension greater than or equal to the 'length' input argument.
    length: The desired length of the TensorArrays.
  Returns:
    tensorarrays: A potentially nested tuple or list of TensorArrays with the
      same structure as 'tensors'. Contains the result of unstacking each Tensor
      in 'tensors'.
  """
  def map_fn(x):
    ta = tf.TensorArray(x.dtype, length, name=x.name.split(':')[0] + '_ta')
    return ta.unstack(x[:length, :])
  return map_nested(map_fn, tensors) 

def testConvertNetworkStateTensorarray(self):
    with self.test_session() as session:
      ta = tf.TensorArray(
          dtype=tf.float32,
          size=0,
          dynamic_size=True,
          clear_after_read=False,
          infer_shape=False)
      # Create a 3-step x 2-stride x 2-feature-dim source array.
      ta = ta.write(0, [[0., 0.]] * 2)  # The zeroth step will be removed.
      ta = ta.write(1, [[1., 10.]] * 2)
      ta = ta.write(2, [[2., 20.]] * 2)
      ta = ta.write(3, [[3., 30.]] * 2)
      tensor = network_units.convert_network_state_tensorarray(ta)
      actual = session.run(tensor)
      self.assertEqual(actual.shape, (6, 2))

      # The arrangement of the values is expected to be stride * steps.
      expected = [[1., 10.], [2., 20.], [3., 30.], [1., 10.], [2., 20.],
                  [3., 30.]]
      self.assertAllEqual(actual, expected) 

def extend_prefixes(prefixes_to_extend, possible_prefix_extensions):
  """Extends prefixes in `prefixes_to_extend` by `possible_prefix_extensions`.

  Args:
    prefixes_to_extend: A 1D tf.string containing prefixes to be extended.
    possible_prefix_extensions: A 1D tf.string containing all possible prefix
      extensions.

  Returns:
    A 1D tf.string containing all the extended prefixes.
  """
  num_new_prefixes = tf.shape(prefixes_to_extend)[0] * tf.shape(
      possible_prefix_extensions)[0]
  extended_prefixes = tf.TensorArray(dtype=tf.string, size=num_new_prefixes)
  position = tf.constant(0, dtype=tf.int32)
  for prefix in prefixes_to_extend:
    for possible_extension in possible_prefix_extensions:
      # [-1] is passed to tf.reshape to flatten the extended prefix. This is
      # important to ensure consistency of shapes.
      extended_prefix = tf.reshape(
          tf.strings.reduce_join([prefix, possible_extension]), [-1])
      extended_prefixes = extended_prefixes.write(position, extended_prefix)
      position += 1
  return extended_prefixes.concat() 

def call(self, inputs):
        x, seq_length = inputs

        x_list = tf.TensorArray(dtype=tf.float32, size=seq_length)
        x_list = x_list.unstack(tf.transpose(x, perm=[1, 0, 2]))
        state = self.cell.get_initial_state(batch_size=self.batch_size, dtype=tf.float32)
        for t in range(seq_length):
            output, state = self.cell(tf.concat([x_list.read(t), tf.zeros([self.batch_size, 1])], axis=1), state)

        output, state = self.cell(self.eof, state)

        output_list = tf.TensorArray(dtype=tf.float32, size=seq_length)
        for t in range(seq_length):
            output, state = self.cell(self.zero, state)
            output_list = output_list.write(t, output[:, 0:self.vector_dim])
        y_pred = tf.sigmoid(tf.transpose(output_list.stack(), perm=[1, 0, 2]))

        return y_pred 

def _maybe_split_batch_beams(self, t, s):
    """Maybe splits the tensor from a batch by beams into a batch of beams.

    We do this so that we can use nest and not run into problems with shapes.

    Args:
      t: `Tensor`, either scalar or shaped `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      If `t` is a matrix or higher order tensor, then the return value is
      `t` reshaped to `[batch_size, beam_width] + s`.  Otherwise `t` is
      returned unchanged.

    Raises:
      ValueError: If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 1:
      return self._split_batch_beams(t, s)
    else:
      return t 

def _maybe_merge_batch_beams(self, t, s):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, `t` is a tensor of dimension `[batch_size * beam_width] + s`,
    then we reshape it to `[batch_size, beam_width] + s`.

    Args:
      t: `Tensor` of dimension `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      A reshaped version of t with shape `[batch_size, beam_width] + s`.

    Raises:
      ValueError:  If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 2:
      return self._merge_batch_beams(t, s)
    else:
      return t 

def _check_batch_beam(t, batch_size, beam_width):
  """Returns an Assert operation checking that the elements of the stacked
  TensorArray can be reshaped to [batch_size, beam_size, -1]. At this point,
  the TensorArray elements have a known rank of at least 1.
  """
  error_message = ("TensorArray reordering expects elements to be "
                   "reshapable to [batch_size, beam_size, -1] which is "
                   "incompatible with the dynamic shape of %s elements. "
                   "Consider setting reorder_tensor_arrays to False to disable "
                   "TensorArray reordering during the beam search."
                   % (t.name))
  rank = t.shape.ndims
  shape = tf.shape(t)
  if rank == 2:
    condition = tf.equal(shape[1], batch_size * beam_width)
  else:
    condition = tf.logical_or(
        tf.equal(shape[1], batch_size * beam_width),
        tf.logical_and(
            tf.equal(shape[1], batch_size),
            tf.equal(shape[2], beam_width)))
  return tf.Assert(condition, [error_message]) 

def _maybe_split_batch_beams(self, t, s):
    """Maybe splits the tensor from a batch by beams into a batch of beams.

    We do this so that we can use nest and not run into problems with shapes.

    Args:
      t: `Tensor`, either scalar or shaped `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      If `t` is a matrix or higher order tensor, then the return value is
      `t` reshaped to `[batch_size, beam_width] + s`.  Otherwise `t` is
      returned unchanged.

    Raises:
      ValueError: If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 1:
      return self._split_batch_beams(t, s)
    else:
      return t 

def _maybe_merge_batch_beams(self, t, s):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, `t` is a tensor of dimension `[batch_size * beam_width] + s`,
    then we reshape it to `[batch_size, beam_width] + s`.

    Args:
      t: `Tensor` of dimension `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      A reshaped version of t with shape `[batch_size, beam_width] + s`.

    Raises:
      ValueError:  If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 2:
      return self._merge_batch_beams(t, s)
    else:
      return t 

def testConvertNetworkStateTensorarray(self):
    with self.test_session() as session:
      ta = tf.TensorArray(
          dtype=tf.float32,
          size=0,
          dynamic_size=True,
          clear_after_read=False,
          infer_shape=False)
      # Create a 3-step x 2-stride x 2-feature-dim source array.
      ta = ta.write(0, [[0., 0.]] * 2)  # The zeroth step will be removed.
      ta = ta.write(1, [[1., 10.]] * 2)
      ta = ta.write(2, [[2., 20.]] * 2)
      ta = ta.write(3, [[3., 30.]] * 2)
      tensor = network_units.convert_network_state_tensorarray(ta)
      actual = session.run(tensor)
      self.assertEqual(actual.shape, (6, 2))

      # The arrangement of the values is expected to be stride * steps.
      expected = [[1., 10.], [2., 20.], [3., 30.], [1., 10.], [2., 20.],
                  [3., 30.]]
      self.assertAllEqual(actual, expected) 

def non_max_suppression(inputs, scores, batch_size, max_output_size,
                        score_threshold=0.7, iou_threshold=0.7, nonempty=False, name='nms'):
    """ Perform NMS on batch of images. """
    with tf.variable_scope(name):
        ix = tf.constant(0)
        filtered_rois = tf.TensorArray(dtype=tf.int32, size=batch_size, infer_shape=False)
        loop_cond = lambda ix, filtered_rois: tf.less(ix, batch_size)
        def _loop_body(ix, filtered_rois):
            indices, score, roi = _filter_tensor(scores[ix], score_threshold, inputs[ix]) # pylint: disable=unbalanced-tuple-unpacking
            roi_corners = tf.concat([roi[:, :2], roi[:, :2]+roi[:, 2:]], axis=-1)
            roi_after_nms = tf.image.non_max_suppression(roi_corners, score, max_output_size, iou_threshold)
            if nonempty:
                is_not_empty = lambda: filtered_rois.write(ix,
                                                           tf.cast(tf.gather(indices, roi_after_nms),
                                                                   dtype=tf.int32))
                is_empty = lambda: filtered_rois.write(ix, tf.constant([[0]]))
                filtered_rois = tf.cond(tf.not_equal(tf.shape(indices)[0], 0), is_not_empty, is_empty)
            else:
                filtered_rois = filtered_rois.write(ix, tf.cast(tf.gather(indices, roi_after_nms), dtype=tf.int32))
            return [ix+1, filtered_rois]
        _, res = tf.while_loop(loop_cond, _loop_body, [ix, filtered_rois])
        res = _array_to_tuple(res, batch_size, [-1, 1])
    return res 

def _maybe_merge_batch_beams(self, t, s):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, `t` is a tensor of dimension `[batch_size * beam_width] + s`,
    then we reshape it to `[batch_size, beam_width] + s`.

    Args:
      t: `Tensor` of dimension `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      A reshaped version of t with shape `[batch_size, beam_width] + s`.

    Raises:
      ValueError:  If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 2:
      return self._merge_batch_beams(t, s)
    else:
      return t 

def _check_batch_beam(t, batch_size, beam_width):
  """Returns an Assert operation checking that the elements of the stacked
  TensorArray can be reshaped to [batch_size, beam_size, -1]. At this point,
  the TensorArray elements have a known rank of at least 1.
  """
  error_message = ("TensorArray reordering expects elements to be "
                   "reshapable to [batch_size, beam_size, -1] which is "
                   "incompatible with the dynamic shape of %s elements. "
                   "Consider setting reorder_tensor_arrays to False to disable "
                   "TensorArray reordering during the beam search."
                   % (t.name))
  rank = t.shape.ndims
  shape = tf.shape(t)
  if rank == 2:
    condition = tf.equal(shape[1], batch_size * beam_width)
  else:
    condition = tf.logical_or(
        tf.equal(shape[1], batch_size * beam_width),
        tf.logical_and(
            tf.equal(shape[1], batch_size),
            tf.equal(shape[2], beam_width)))
  return tf.Assert(condition, [error_message]) 

def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
  """Raises an exception if dimensions are known statically and can not be
  reshaped to [batch_size, beam_size, -1].
  """
  reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
  if (batch_size is not None and shape[0].value is not None
      and (shape[0] != batch_size * beam_width
           or (shape.ndims >= 2 and shape[1].value is not None
               and (shape[0] != batch_size or shape[1] != beam_width)))):
    tf.logging.warn("TensorArray reordering expects elements to be "
                    "reshapable to %s which is incompatible with the "
                    "current shape %s. Consider setting "
                    "reorder_tensor_arrays to False to disable TensorArray "
                    "reordering during the beam search."
                    % (reshaped_shape, shape))
    return False
  return True 

def _maybe_split_batch_beams(self, t, s):
    """Maybe splits the tensor from a batch by beams into a batch of beams.

    We do this so that we can use nest and not run into problems with shapes.

    Args:
      t: `Tensor`, either scalar or shaped `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      If `t` is a matrix or higher order tensor, then the return value is
      `t` reshaped to `[batch_size, beam_width] + s`.  Otherwise `t` is
      returned unchanged.

    Raises:
      ValueError: If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 1:
      return self._split_batch_beams(t, s)
    else:
      return t 

def _maybe_merge_batch_beams(self, t, s):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, `t` is a tensor of dimension `[batch_size * beam_width] + s`,
    then we reshape it to `[batch_size, beam_width] + s`.

    Args:
      t: `Tensor` of dimension `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      A reshaped version of t with shape `[batch_size, beam_width] + s`.

    Raises:
      ValueError:  If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 2:
      return self._merge_batch_beams(t, s)
    else:
      return t 

def _maybe_split_batch_beams(self, t, s):
    """Maybe splits the tensor from a batch by beams into a batch of beams.

    We do this so that we can use nest and not run into problems with shapes.

    Args:
      t: `Tensor`, either scalar or shaped `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      If `t` is a matrix or higher order tensor, then the return value is
      `t` reshaped to `[batch_size, beam_width] + s`.  Otherwise `t` is
      returned unchanged.

    Raises:
      ValueError: If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 1:
      return self._split_batch_beams(t, s)
    else:
      return t 

def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
  """Raises an exception if dimensions are known statically and can not be
  reshaped to [batch_size, beam_size, -1].
  """
  reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
  if (batch_size is not None and shape[0].value is not None
      and (shape[0] != batch_size * beam_width
           or (shape.ndims >= 2 and shape[1].value is not None
               and (shape[0] != batch_size or shape[1] != beam_width)))):
    tf.logging.warn("TensorArray reordering expects elements to be "
                    "reshapable to %s which is incompatible with the "
                    "current shape %s. Consider setting "
                    "reorder_tensor_arrays to False to disable TensorArray "
                    "reordering during the beam search."
                    % (reshaped_shape, shape))
    return False
  return True 

def _maybe_merge_batch_beams(self, t, s):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, `t` is a tensor of dimension `[batch_size * beam_width] + s`,
    then we reshape it to `[batch_size, beam_width] + s`.

    Args:
      t: `Tensor` of dimension `[batch_size * beam_width] + s`.
      s: `Tensor`, Python int, or `TensorShape`.

    Returns:
      A reshaped version of t with shape `[batch_size, beam_width] + s`.

    Raises:
      ValueError:  If the rank of `t` is not statically known.
    """
    if isinstance(t, tf.TensorArray):
      return t
    _check_maybe(t)
    if t.shape.ndims >= 2:
      return self._merge_batch_beams(t, s)
    else:
      return t 

def _compute_states(self):

    _inputs = tf.transpose(self.inputs, [1, 0, 2])
    x_ta = tf.TensorArray(tf.float32, size=self.length).unstack(_inputs)
    h_ta = tf.TensorArray(tf.float32, size=self.length)
    c_ta = tf.TensorArray(tf.float32, size=self.length)

    def cond(t, c, h, c_ta, h_ta):
      return tf.less(t, self.length)

    def body(t, c, h, c_ta, h_ta):

      x = x_ta.read(t)
      num_units, input_size = self.num_hidden_units, self.input_size

      with tf.variable_scope('lstm'):
        c_tilde = self.activation(self._linear(h, x, num_units, scope='c'))
        i = tf.nn.sigmoid(self._linear(h, x, num_units, scope='i'))
        f = tf.nn.sigmoid(self._linear(h, x, num_units, shift=self.optional_bias_shift, scope='f'))
        o = tf.nn.sigmoid(self._linear(h, x, num_units, scope='o'))
        c_new = i * c_tilde + f * c
        h_new = o * self.activation(c_new)

      c_ta_new = c_ta.write(t, c_new)
      h_ta_new = h_ta.write(t, h_new)
      return t + 1, c_new, h_new, c_ta_new, h_ta_new

    t = tf.constant(0)
    c, h = tf.split(tf.squeeze(self.initial_states, [1]), 2, axis=1)
    _, _, _, c_ta, h_ta = tf.while_loop(cond, body, [t, c, h, c_ta, h_ta])

    outputs = tf.transpose(h_ta.stack(), [1, 0, 2], name='outputs')
    cells = tf.transpose(c_ta.stack(), [1, 0, 2])
    states = tf.concat([cells, outputs], axis=2, name='states')
    return outputs, states 

def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size,
                                range_size, gather_shape):
  """Maybe applies _tensor_gather_helper.

  This applies _tensor_gather_helper when the gather_from dims is at least as
  big as the length of gather_shape. This is used in conjunction with nest so
  that we don't apply _tensor_gather_helper to inapplicable values like scalars.

  Args:
    gather_indices: The tensor indices that we use to gather.
    gather_from: The tensor that we are gathering from.
    batch_size: The batch size.
    range_size: The number of values in each range. Likely equal to beam_width.
    gather_shape: What we should reshape gather_from to in order to preserve the
      correct values. An example is when gather_from is the attention from an
      AttentionWrapperState with shape [batch_size, beam_width, attention_size].
      There, we want to preserve the attention_size elements, so gather_shape is
      [batch_size * beam_width, -1]. Then, upon reshape, we still have the
      attention_size as desired.

  Returns:
    output: Gathered tensor of shape tf.shape(gather_from)[:1+len(gather_shape)]
      or the original tensor if its dimensions are too small.
  """
  if isinstance(gather_from, tf.TensorArray):
    return gather_from
  _check_maybe(gather_from)
  if gather_from.shape.ndims >= len(gather_shape):
    return _tensor_gather_helper(
        gather_indices=gather_indices,
        gather_from=gather_from,
        batch_size=batch_size,
        range_size=range_size,
        gather_shape=gather_shape)
  else:
    return gather_from 

def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size,
                                range_size, gather_shape):
  """Maybe applies _tensor_gather_helper.

  This applies _tensor_gather_helper when the gather_from dims is at least as
  big as the length of gather_shape. This is used in conjunction with nest so
  that we don't apply _tensor_gather_helper to inapplicable values like scalars.

  Args:
    gather_indices: The tensor indices that we use to gather.
    gather_from: The tensor that we are gathering from.
    batch_size: The batch size.
    range_size: The number of values in each range. Likely equal to beam_width.
    gather_shape: What we should reshape gather_from to in order to preserve the
      correct values. An example is when gather_from is the attention from an
      AttentionWrapperState with shape [batch_size, beam_width, attention_size].
      There, we want to preserve the attention_size elements, so gather_shape is
      [batch_size * beam_width, -1]. Then, upon reshape, we still have the
      attention_size as desired.

  Returns:
    output: Gathered tensor of shape tf.shape(gather_from)[:1+len(gather_shape)]
      or the original tensor if its dimensions are too small.
  """
  if isinstance(gather_from, tf.TensorArray):
    return gather_from
  _check_maybe(gather_from)
  if gather_from.shape.ndims >= len(gather_shape):
    return _tensor_gather_helper(
        gather_indices=gather_indices,
        gather_from=gather_from,
        batch_size=batch_size,
        range_size=range_size,
        gather_shape=gather_shape)
  else:
    return gather_from 

def _maybe_sort_array_beams(self, t, parent_ids, sequence_length):
    """Maybe sorts beams within a `TensorArray`.

    Args:
      t: A `TensorArray` of size `max_time` that contains `Tensor`s of shape
        `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]` where
        `s` is the depth shape.
      parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
      sequence_length: The sequence length of shape `[batch_size, beam_width]`.

    Returns:
      A `TensorArray` where beams are sorted in each `Tensor` or `t` itself if
      it is not a `TensorArray` or does not meet shape requirements.
    """
    if not isinstance(t, tf.TensorArray):
      return t
    # pylint: disable=protected-access
    if (not t._infer_shape or not t._element_shape
        or t._element_shape[0].ndims is None
        or t._element_shape[0].ndims < 1):
      shape = (
          t._element_shape[0] if t._infer_shape and t._element_shape
          else tf.TensorShape(None))
      tf.logging.warn("The TensorArray %s in the cell state is not amenable to "
                      "sorting based on the beam search result. For a "
                      "TensorArray to be sorted, its elements shape must be "
                      "defined and have at least a rank of 1, but saw shape: %s"
                      % (t.handle.name, shape))
      return t
    shape = t._element_shape[0]
    # pylint: enable=protected-access
    if not _check_static_batch_beam_maybe(
        shape, tf.contrib.util.constant_value(self._batch_size),
        self._beam_width):
      return t
    t = t.stack()
    with tf.control_dependencies(
        [_check_batch_beam(t, self._batch_size, self._beam_width)]):
      return gather_tree_from_array(t, parent_ids, sequence_length) 

def _compute_states(self):
    """ Compute hidden states.

    Returns:
      A tuple, (outputs, states).
    """

    _inputs = tf.transpose(self.inputs, [1, 0, 2])
    x_ta = tf.TensorArray(tf.float32, size=self.length).unstack(_inputs)
    h_ta = tf.TensorArray(tf.float32, size=self.length)

    def cond(t, h, h_ta):
      return tf.less(t, self.length)

    def body(t, h, h_ta):

      x = x_ta.read(t)
      num_units, input_size = self.num_hidden_units, self.input_size

      with tf.variable_scope('simple_rnn'):
        h_new = self.activation(self._linear(h, x, num_units, scope='simple_rnn'))

      h_ta_new = h_ta.write(t, h_new)
      return t + 1, h_new, h_ta_new

    t = tf.constant(0)
    h = tf.squeeze(self.initial_states, [1])
    _, _, h_ta = tf.while_loop(cond, body, [t, h, h_ta])

    states = tf.transpose(h_ta.stack(), [1, 0, 2], name='states')
    outputs = tf.identity(states, name='outputs')
    return outputs, states 

def _initialize(self, batch_size, start_ids, attention_size=None):
    start_ids = tfa.seq2seq.tile_batch(start_ids, self.beam_size)
    finished = tf.zeros([batch_size * self.beam_size], dtype=tf.bool)
    # Give all probability to first beam for the first iteration.
    initial_log_probs = tf.tile([0.] + [-float("inf")] * (self.beam_size - 1), [batch_size])
    parent_ids = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
    sequence_lengths = tf.zeros([batch_size * self.beam_size], dtype=tf.int32)
    extra_vars = [parent_ids, sequence_lengths]
    if self.coverage_penalty != 0:
      if attention_size is None:
        raise ValueError("The attention size should be known to support coverage penalty")
      accumulated_attention = tf.zeros([batch_size * self.beam_size, attention_size])
      extra_vars.append(accumulated_attention)
    return start_ids, finished, initial_log_probs, tuple(extra_vars) 

def _get_shape_invariants(tensor):
  """Returns the shape of the tensor but sets middle dims to None."""
  if isinstance(tensor, tf.TensorArray):
    shape = None
  else:
    shape = tensor.shape.as_list()
    for i in range(1, len(shape) - 1):
      shape[i] = None
  return tf.TensorShape(shape)