content: from collections import namedtuple
  import tensorflow as tf
  
  
  class TimedNTM(object):
      """Timed Neural Turing Machine
  
      Inspired by paper:
          https://arxiv.org/pdf/1410.5401.pdf
      Implementation inspired by:
          https://github.com/carpedm20/NTM-tensorflow/blob/master/ntm_cell.py
  
      See our paper for details: https://arxiv.org/abs/1811.11707
      """
  
      def __init__(self, attn_shift_range, sparse_attention, name):
          """Construct the `TimedNTM`.
  
          Args:
              attn_shift_range: Python int.
                  A time range within which to attend to the memory by location
              sparse_attention: Python bool.
                  If `True` use sparsemax instead of softmax for probs
              name: Name to use when creating ops.
          """
  
          # interpolation gate
          self.name = 'timed_ntm_' + name
  
          self._inter_gate = tf.layers.Dense(
              units=1,
              activation=tf.sigmoid,
              name=self.name + '/inter_gate'
          )
          # if use sparsemax instead of softmax for probs
          self._sparse_attention = sparse_attention
  
          if sparse_attention:
              # sparsemax doesn't support inf
              self._inf = float(5000)
          else:
              self._inf = float('inf')
  
          # shift weighting if range is provided
          if attn_shift_range:
              self._shift_weight = tf.layers.Dense(
                  units=2 * attn_shift_range + 1,
                  activation=tf.nn.softmax,
                  name=self.name + '/shift_weight'
              )
          else:
              self._shift_weight = None
  
          # sharpening parameter
          self._gamma_sharp = tf.layers.Dense(
              units=1,
              activation=lambda a: tf.nn.softplus(a) + 1,
              bias_initializer=tf.constant_initializer(1),
              name=self.name + '/gamma_sharp'
          )
  
      def __call__(self, attn_inputs, scores, scores_state, mask):
          # apply exponential moving average with interpolation gate weight
          # to scores from previous time which are equal to probs at this point
          # different from original NTM where it is applied after softmax
          i_g = self._inter_gate(attn_inputs)
  
          # scores limited by time
          scores = tf.concat([i_g * scores[:, :-1] + (1 - i_g) * scores_state,
                              scores[:, -1:]], 1)
          next_scores_state = scores
  
          if mask is not None:
              # apply mask to scores
              if self._shift_weight is not None:
                  # rearrange scores to make them continuous for convolution
                  scores = tf.map_fn(self._rearrange_fn,
                                     [scores, mask], dtype=scores.dtype)
              else:
                  scores = tf.where(mask > 0,
                                    scores, -self._inf * tf.ones_like(scores))
  
          # create probabilities for attention
          if self._sparse_attention:
              probs = tf.contrib.sparsemax.sparsemax(scores)
          else:
              probs = tf.nn.softmax(scores)
  
          if self._shift_weight is not None:
              s_w = self._shift_weight(attn_inputs)
  
              # we want to go back in time during convolution
              conv_probs = tf.reverse(probs, axis=[1])
  
              # preare probs for tf.nn.depthwise_conv2d
              # [in_width, in_channels=batch]
              conv_probs = tf.transpose(conv_probs, [1, 0])
              # [batch=1, in_height=1, in_width=time+1, in_channels=batch]
              conv_probs = conv_probs[tf.newaxis, tf.newaxis, :, :]
  
              # [filter_height=1, filter_width=2*attn_shift_range+1,
              #   in_channels=batch, channel_multiplier=1]
              conv_s_w = tf.transpose(s_w, [1, 0])
              conv_s_w = conv_s_w[tf.newaxis, :, :, tf.newaxis]
  
              # perform 1d convolution
              # [batch=1, out_height=1, out_width=time+1, out_channels=batch]
              conv_probs = tf.nn.depthwise_conv2d_native(conv_probs, conv_s_w,
                                                         [1, 1, 1, 1], 'SAME')
              conv_probs = conv_probs[0, 0, :, :]
              conv_probs = tf.transpose(conv_probs, [1, 0])
  
              probs = tf.reverse(conv_probs, axis=[1])
  
              if mask is not None:
                  # arrange probs back to their original time order
                  probs = tf.map_fn(self._arrange_back_fn,
                                    [probs, mask], dtype=probs.dtype)
  
          # sharpening
          g_sh = self._gamma_sharp(attn_inputs)
  
          powed_probs = tf.pow(probs, g_sh)
          probs = powed_probs / (
              tf.reduce_sum(powed_probs, 1, keepdims=True) + 1e-32)
  
          return probs, next_scores_state
  
      def _rearrange_fn(self, list_tensor_1d_mask_1d):
          """Rearranges tensor_1d to put all the values
              where mask_1d=1 to the right and
              where mask_1d=0 to the left and sets them to -infinity"""
          tensor_1d, mask_1d = list_tensor_1d_mask_1d
  
          partitioned_tensor = tf.dynamic_partition(tensor_1d,
                                                    mask_1d, 2)
          partitioned_tensor[0] = \
              -self._inf * tf.ones_like(partitioned_tensor[0])
  
          return tf.concat(partitioned_tensor, 0)
  
      @staticmethod
      def _arrange_back_fn(list_tensor_1d_mask_1d):
          """Arranges back tensor_1d to restore original order
              modified by `_rearrange_fn` according to mask_1d:
              - number of 0s in mask_1d values on the left are set to
                their corresponding places where mask_1d=0,
              - number of 1s in mask_1d values on the right are set to
                their corresponding places where mask_1d=1"""
          tensor_1d, mask_1d = list_tensor_1d_mask_1d
  
          mask_indices = tf.dynamic_partition(tf.range(tf.shape(tensor_1d)[0]),
                                              mask_1d, 2)
  
          mask_sum = tf.reduce_sum(mask_1d, axis=0)
          partitioned_tensor = [tf.zeros_like(tensor_1d[:-mask_sum]),
                                tensor_1d[-mask_sum:]]
  
          return tf.dynamic_stitch(mask_indices, partitioned_tensor)
  
  
  def _compute_time_attention(attention_mechanism, attn_inputs, attention_state,
                              # time is added to calculate time attention
                              time, timed_ntm, time_mask, ignore_mask,
                              attention_layer):
      """Computes the attention and alignments limited by time
          for a given attention_mechanism.
  
          Modified helper method from tensorflow."""
  
      scores, _ = attention_mechanism(attn_inputs, state=attention_state)
  
      # take only scores from current and past times
      timed_scores = scores[:, :time + 1]
      timed_scores_state = attention_state[:, :time]
  
      # get mask for past times
      timed_time_mask = time_mask[:, :time]
      if ignore_mask is not None:
          timed_time_mask *= 1 - ignore_mask[:, :time]
  
      # set mask for current time to 1
      timed_time_mask = tf.concat([timed_time_mask,
                                   tf.ones_like(time_mask[:, :1])], 1)
  
      # pass these scores to NTM
      probs, next_scores_state = timed_ntm(attn_inputs, timed_scores,
                                           timed_scores_state,
                                           timed_time_mask)
  
      # concatenate probs with zeros to get new alignments
      zeros = tf.zeros_like(scores)
      # remove current time from attention
      alignments = tf.concat([probs[:, :-1], zeros[:, time:]], 1)
  
      # Reshape from [batch_size, memory_time] to [batch_size, 1, memory_time]
      expanded_alignments = tf.expand_dims(alignments, 1)
  
      # Context is the inner product of alignments and values along the
      # memory time dimension.
      # alignments shape is
      #   [batch_size, 1, memory_time]
      # attention_mechanism.values shape is
      #   [batch_size, memory_time, memory_size]
      # the batched matmul is over memory_time, so the output shape is
      #   [batch_size, 1, memory_size].
      # we then squeeze out the singleton dim.
      context = tf.matmul(expanded_alignments, attention_mechanism.values)
      context = tf.squeeze(context, [1])
  
      if attention_layer is not None:
          attention = attention_layer(tf.concat([attn_inputs, context], 1))
      else:
          attention = context
  
      # return current time to attention
      alignments = tf.concat([probs, zeros[:, time + 1:]], 1)
      next_attention_state = tf.concat([next_scores_state,
                                        zeros[:, time + 1:]], 1)
      return attention, alignments, next_attention_state
  
  
  # noinspection PyProtectedMember
  class TimeAttentionWrapperState(
      namedtuple("TimeAttentionWrapperState",
                 tf.contrib.seq2seq.AttentionWrapperState._fields +
                 ("all_time_masks", "all_cell_states"))):  # added
      """Modified  from tensorflow's tf.contrib.seq2seq.AttentionWrapperState
          see there for description of the parameters
  
      Additional fields:
          - `all_time_masks`: A mask applied to a memory
             that filters certain time steps
          - `all_cell_states`: All states of the wrapped `RNNCell`
             at all the previous time steps.
      """
  
      def clone(self, **kwargs):
          """Copied  from tensorflow's tf.contrib.seq2seq.AttentionWrapperState
              see there for description of the parameters"""
  
          def with_same_shape(old, new):
              """Check and set new tensor's shape."""
              if isinstance(old, tf.Tensor) and isinstance(new, tf.Tensor):
                  return tf.contrib.framework.with_same_shape(old, new)
              return new
  
          return tf.contrib.framework.nest.map_structure(
              with_same_shape,
              self,
              super(TimeAttentionWrapperState, self)._replace(**kwargs)
          )
  
  
  class TimeAttentionWrapper(tf.contrib.seq2seq.AttentionWrapper):
      """Custom AttentionWrapper that takes into account time
          when calculating attention.
          Attention is calculated before calling rnn cell.
  
          Modified from tensorflow's tf.contrib.seq2seq.AttentionWrapper.
  
          See our paper for details: https://arxiv.org/abs/1811.11707
      """
  
      def __init__(self, cell,
                   attention_mechanism,
                   sequence_len,
                   attn_shift_range=0,
                   sparse_attention=False,
                   attention_layer_size=None,
                   alignment_history=False,
                   rnn_and_attn_inputs_fn=None,
                   ignore_mask=None,
                   cell_input_fn=None,
                   index_of_attn_to_copy=None,
                   likelihood_fn=None,
                   tensor_not_to_copy=None,
                   output_attention=False,
                   initial_cell_state=None,
                   name=None,
                   attention_layer=None):
          """Construct the `TimeAttentionWrapper`.
              See the super class for the original arguments description.
  
          Additional args:
              sequence_len: Python integer.
                  Maximum length of the sequence, used to create
                  appropriate TensorArray for all cell states
                  in TimeAttentionWrapperState
              attn_shift_range: Python integer (`0` by default).
                  A time range within which to attend to the memory
                  by location in Neural Turing Machine.
              sparse_attention: Python bool.
                  A flag to use sparsemax (if `True`) instead of
                  softmax (if `False`, default) for probabilities
              inputs_and_attn_inputs_fn: (optional) A `callable`.
                  A function that creates inputs and attention inputs tensors.
              ignore_mask: (optional) Boolean Tensor.
                  Determines which time steps to ignore in attention
              index_of_attn_to_copy: (optional) Python integer.
                  An index of attention mechanism that picks
                  which part of attention tensor to use for copying to output,
                  the default is `None`, which turns off copying mechanism.
                  Copy inspired by: https://arxiv.org/pdf/1603.06393.pdf
              likelihood_fn: (optional) A `callable`.
                  A method to perform likelihood calculation to
                  filter time step in copy mechanism.
                  Returns a tuple of binary likelihood and likelihood
              tensor_not_to_copy: (optional) A Tensor.
                  A tensor, which shouldn't be copied from previous time steps
  
          Modified args:
              output_attention: Python bool.  If `True`, the output at each
                  time step is the concatenated cell outputs,
                  attention values and additional values described in
                  `additional_output_size()`, used in copy mechanism.
          """
          super(TimeAttentionWrapper, self).__init__(
              cell,
              attention_mechanism,
              attention_layer_size,
              alignment_history,
              cell_input_fn,
              output_attention,
              initial_cell_state,
              name,
              attention_layer
          )
          self._sequence_len = sequence_len
  
          if not isinstance(attn_shift_range, list):
              # attn_shift_range might not be a list
              attn_shift_range = [attn_shift_range]
          self._timed_ntms = [TimedNTM(attn_shift_range[0],
                                       sparse_attention,
                                       name='0')]
          if self._is_multi:
              # if there are several attention mechanisms,
              # create additional TimedNTMs for them
              if len(attn_shift_range) == 1:
                  # original attn_shift_range might not be a list
                  attn_shift_range *= len(attention_mechanism)
              elif len(attn_shift_range) != len(attention_mechanism):
                  raise ValueError(
                      "If provided, `attn_shift_range` must contain exactly one "
                      "integer per attention_mechanism, saw: {} vs {}"
                      "".format(len(attn_shift_range), len(attention_mechanism))
                  )
              for i in range(1, len(attention_mechanism)):
                  self._timed_ntms.append(TimedNTM(attn_shift_range[i],
                                                   sparse_attention,
                                                   name=str(i)))
  
          if rnn_and_attn_inputs_fn is None:
              rnn_and_attn_inputs_fn = self._default_rnn_and_attn_inputs_fn
          else:
              if not callable(rnn_and_attn_inputs_fn):
                  raise TypeError(
                      "`rnn_and_attn_inputs_fn` must be callable, saw type: {}"
                      "".format(type(rnn_and_attn_inputs_fn).__name__)
                  )
          self._rnn_and_attn_inputs_fn = rnn_and_attn_inputs_fn
  
          if not isinstance(ignore_mask, list):
              self._ignore_mask = [tf.cast(ignore_mask, tf.int32)]
          else:
              self._ignore_mask = [tf.cast(i_m, tf.int32) for i_m in ignore_mask]
  
          self._index_of_attn_to_copy = index_of_attn_to_copy
  
          self._likelihood_fn = likelihood_fn
          self._tensor_not_to_copy = tensor_not_to_copy
  
      @staticmethod
      def _default_rnn_and_attn_inputs_fn(inputs, cell_state):
          if isinstance(cell_state, tf.contrib.rnn.LSTMStateTuple):
              return inputs, tf.concat([inputs, cell_state.h], -1)
          else:
              return inputs, tf.concat([inputs, cell_state], -1)
  
      @staticmethod
      def additional_output_size():
          """Number of additional outputs:
  
          likelihoods:
              attn_likelihood, state_likelihood
          debugging info:
              current_time_prob,
              bin_likelihood_not_to_copy, bin_likelihood_to_copy
  
          **Method should be static**
          """
          return 2 + 3
  
      @property
      def output_size(self):
          if self._output_attention:
              if self._index_of_attn_to_copy is not None:
                  # output both raw rnn cell_output and
                  # cell_output with copied attention
                  # together with attention vector itself
                  # and additional output
                  return (2 * self._cell.output_size +
                          self._attention_layer_size +
                          self.additional_output_size())
              else:
                  return self._cell.output_size + self._attention_layer_size
          else:
              return self._cell.output_size
  
      @property
      def state_size(self):
          """The `state_size` property of `TimeAttentionWrapper`.
          Returns:
              A `TimeAttentionWrapperState` tuple containing shapes
              used by this object.
          """
  
          # use AttentionWrapperState from superclass
          state_size = super(TimeAttentionWrapper, self).state_size
  
          all_cell_states = self._cell.state_size
  
          return TimeAttentionWrapperState(
              cell_state=state_size.cell_state,
              time=state_size.time,
              attention=state_size.attention,
              alignments=state_size.alignments,
              attention_state=state_size.attention_state,
              alignment_history=state_size.alignment_history,
              all_time_masks=self._sequence_len,
              all_cell_states=all_cell_states)
  
      def zero_state(self, batch_size, dtype):
          """Modified  from tensorflow's zero_state
              see there for description of the parameters"""
  
          # use AttentionWrapperState from superclass
          zero_state = super(TimeAttentionWrapper,
                             self).zero_state(batch_size, dtype)
  
          with tf.name_scope(type(self).__name__ + "ZeroState",
                             values=[batch_size]):
              # store time masks
              all_time_masks = tf.TensorArray(
                  tf.int32,
                  size=self._sequence_len + 1,
                  dynamic_size=False,
                  clear_after_read=False
              ).write(0, tf.zeros([batch_size, self.state_size.all_time_masks],
                                  tf.int32))
  
              # store all cell states into a tensor array to allow
              # copy mechanism to go back in time
              if isinstance(self._cell.state_size,
                            tf.contrib.rnn.LSTMStateTuple):
                  all_cell_states = tf.contrib.rnn.LSTMStateTuple(
                      tf.TensorArray(dtype, size=self._sequence_len + 1,
                                     dynamic_size=False,
                                     clear_after_read=False
                                     ).write(0, zero_state.cell_state.c),
                      tf.TensorArray(dtype, size=self._sequence_len + 1,
                                     dynamic_size=False,
                                     clear_after_read=False
                                     ).write(0, zero_state.cell_state.h)
                  )
              else:
                  all_cell_states = tf.TensorArray(
                      dtype, size=0,
                      dynamic_size=False,
                      clear_after_read=False
                  ).write(0, zero_state.cell_state)
  
              return TimeAttentionWrapperState(
                  cell_state=zero_state.cell_state,
                  time=zero_state.time,
                  attention=zero_state.attention,
                  alignments=zero_state.alignments,
                  attention_state=zero_state.attention_state,
                  alignment_history=zero_state.alignment_history,
                  all_time_masks=all_time_masks,
                  all_cell_states=all_cell_states
              )
  
      def call(self, inputs, state):
          """Perform a step of attention-wrapped RNN.
  
          The order has changed:
          - Step 1: Calculate attention inputs based on the previous cell state
                    and current inputs
          - Step 2: Score the output with `attention_mechanism`.
          - Step 3: Calculate the alignments by passing the score through the
                    `normalizer` and limit them by time.
          - Step 4: Calculate the context vector as the inner product between the
                    alignments and the attention_mechanism's values (memory).
          - Step 5: Calculate the attention output by concatenating
                    the cell output and context through the attention layer
                    (a linear layer with `attention_layer_size` outputs).
          - Step 6: Mix the `inputs` and `attention` output via
                    `cell_input_fn` to get cell inputs.
          - Step 7: Call the wrapped `cell` with these cell inputs and
                    its previous state.
          - Step 8: (optional) Maybe copy output and cell state from history
  
          Args:
            inputs: (Possibly nested tuple of) Tensor,
                    the input at this time step.
            state: An instance of `TimeAttentionWrapperState`
                   containing tensors from the previous time step.
  
          Returns:
            A tuple `(attention_or_cell_output, next_state)`, where:
  
            - `attention_or_cell_output` depending on `output_attention`.
            - `next_state` is an instance of `TimeAttentionWrapperState`
               containing the state calculated at this time step.
  
          Raises:
            TypeError: If `state` is not an instance of
            `TimeAttentionWrapperState`.
          """
          if not isinstance(state, TimeAttentionWrapperState):
              raise TypeError("Expected state to be instance of "
                              "TimeAttentionWrapperState. "
                              "Received type {} instead.".format(type(state)))
  
          # Step 1: Calculate attention based on
          #         the previous output and current input
          cell_state = state.cell_state
  
          rnn_inputs, attn_inputs = self._rnn_and_attn_inputs_fn(inputs,
                                                                 cell_state)
  
          cell_batch_size = (
              attn_inputs.shape[0].value or
              tf.shape(attn_inputs)[0])
          error_message = (
              "When applying AttentionWrapper %s: " % self.name +
              "Non-matching batch sizes between the memory "
              "(encoder output) and the query (decoder output).  "
              "Are you using "
              "the BeamSearchDecoder?  "
              "You may need to tile your memory input via "
              "the tf.contrib.seq2seq.tile_batch function with argument "
              "multiple=beam_width.")
          with tf.control_dependencies(
                  self._batch_size_checks(cell_batch_size, error_message)):
              attn_inputs = tf.identity(
                  attn_inputs, name="checked_attn_inputs")
  
          if self._is_multi:
              previous_attention_state = state.attention_state
              previous_alignment_history = state.alignment_history
          else:
              previous_attention_state = [state.attention_state]
              previous_alignment_history = [state.alignment_history]
  
          all_alignments = []
          all_attentions = []
          all_attention_states = []
          maybe_all_histories = []
  
          prev_time_masks = self._read_from_tensor_array(state.all_time_masks,
                                                         state.time)
          prev_time_mask = prev_time_masks[:, -1, :]
  
          for i, attention_mechanism in enumerate(self._attention_mechanisms):
              # Steps 2 - 5 are performed inside `_compute_time_attention`
              (attention, alignments,
               next_attention_state) = _compute_time_attention(
                  attention_mechanism, attn_inputs,
                  previous_attention_state[i],
                  # time is added to calculate time attention
                  state.time, self._timed_ntms[i],
                  # provide boolean masks, to ignore some time steps
                  prev_time_mask, self._ignore_mask[i],
                  self._attention_layers[i]
                  if self._attention_layers else None)
  
              alignment_history = previous_alignment_history[i].write(
                  state.time, alignments) if self._alignment_history else ()
  
              all_attention_states.append(next_attention_state)
              all_alignments.append(alignments)
              all_attentions.append(attention)
              maybe_all_histories.append(alignment_history)
  
          attention = tf.concat(all_attentions, 1)
  
          # Step 6: Mix the `inputs` and `attention` output via
          #         `cell_input_fn` to get cell inputs.
          cell_inputs = self._cell_input_fn(rnn_inputs, attention)
  
          # Step 7: Call the wrapped `cell` with these cell inputs and
          #         its previous state.
          cell_output, next_cell_state = self._cell(cell_inputs, cell_state)
  
          prev_all_cell_states = state.all_cell_states
  
          time_mask = tf.concat([prev_time_mask[:, :state.time],
                                 tf.ones_like(prev_time_mask[:, :1]),
                                 prev_time_mask[:, state.time + 1:]], 1)
  
          if self._index_of_attn_to_copy is not None:
              # Step 8: Maybe copy output and cell state from history
  
              # get relevant previous outputs from history
              attn_to_copy = all_attentions[self._index_of_attn_to_copy]
              # copy them to current output
              cell_output_with_attn = cell_output + attn_to_copy
  
              memory_probs = self._get_memory_probs(all_alignments, state.time)
  
              # check that we do not pay attention to `tensor_not_to_copy`
              bin_likelihood_not_to_copy, _ = self._likelihood_fn(
                  cell_output_with_attn, self._tensor_not_to_copy)
              # recalculate probs
              memory_probs *= 1 - bin_likelihood_not_to_copy
  
              history_alignments = self._history_alignments(memory_probs)
  
              # get previous output from the history
              prev_output = self._prev_output(cell_output_with_attn,
                                              history_alignments,
                                              state.time)
  
              # check that current output is close to
              # the one in the history to which we pay attention to
              bin_likelihood_to_copy, _ = self._likelihood_fn(
                  cell_output_with_attn, prev_output)
              # recalculate probs
              memory_probs *= bin_likelihood_to_copy
  
              history_alignments = self._history_alignments(memory_probs)
              current_time_prob = history_alignments[:, -1:]
  
              # create additional likelihoods to maximize
              attn_likelihood = self._additional_likelihood(
                  attn_to_copy,
                  prev_output,
                  current_time_prob
              )
              state_likelihood = self._additional_likelihood(
                  cell_output + tf.stop_gradient(attn_to_copy),
                  prev_output,
                  current_time_prob
              )
  
              # recalculate time_mask
              time_mask = self._apply_alignments_to_history(
                  tf.cast(history_alignments, time_mask.dtype),
                  prev_time_masks[:, :-1, :],
                  time_mask
              )
  
              # recalculate new next_cell_state based on history_alignments
              next_cell_state = self._new_next_cell_state(
                  prev_all_cell_states,
                  next_cell_state,
                  cell_output_with_attn,
                  history_alignments,
                  state.time
              )
  
              all_cell_states = self._all_cell_states(
                  prev_all_cell_states,
                  next_cell_state,
                  state.time
              )
  
              if self._output_attention:
                  # concatenate cell outputs, attention, additional likelihoods
                  # and copy_attn_debug
                  output = tf.concat([cell_output_with_attn,
                                      cell_output,
                                      attention,
                                      # additional likelihoods
                                      attn_likelihood, state_likelihood,
                                      # copy_attn_debug
                                      bin_likelihood_not_to_copy,
                                      bin_likelihood_to_copy,
                                      current_time_prob], 1)
              else:
                  output = cell_output_with_attn
  
          else:
              # do not waste resources on storing history
              all_cell_states = prev_all_cell_states
  
              if self._output_attention:
                  output = tf.concat([cell_output, attention], 1)
              else:
                  output = cell_output
  
          all_time_masks = state.all_time_masks.write(state.time + 1, time_mask)
  
          next_state = TimeAttentionWrapperState(
              time=state.time + 1,
              cell_state=next_cell_state,
              attention=attention,
              attention_state=self._item_or_tuple(all_attention_states),
              alignments=self._item_or_tuple(all_alignments),
              alignment_history=self._item_or_tuple(maybe_all_histories),
              all_time_masks=all_time_masks,
              all_cell_states=all_cell_states
          )
          return output, next_state
  
      # helper for TensorArray
      @staticmethod
      def _read_from_tensor_array(tensor_array, time):
          """TensorArray time reader"""
          return tf.transpose(tensor_array.gather(tf.range(0, time + 1)),
                              [1, 0, 2])
  
      # helper methods for copy mechanism
      def _get_memory_probs(self, all_alignments, time):
          """Helper method to get memory_probs from all_alignments"""
  
          memory_probs = tf.stop_gradient(all_alignments[
                                              self._index_of_attn_to_copy][:,
                                          :time])
  
          # binarize memory_probs only if max value is larger than margin=0.1
          memory_probs_max = tf.reduce_max(memory_probs, axis=1, keepdims=True)
          memory_probs_max = tf.where(memory_probs_max > 0.1,
                                      memory_probs_max, -memory_probs_max)
  
          return tf.where(tf.equal(memory_probs, memory_probs_max),
                          tf.ones_like(memory_probs),
                          tf.zeros_like(memory_probs))
  
      @staticmethod
      def _history_alignments(memory_probs):
          """Helper method to apply binary mask to memory_probs"""
  
          current_time_prob = 1 - tf.reduce_sum(memory_probs, 1, keepdims=True)
          return tf.concat([memory_probs, current_time_prob], 1)
  
      @staticmethod
      def _apply_alignments_to_history(alignments, history_states, state):
          """Helper method to apply attention probabilities to rnn history
  
          copied from tf's `_compute_attention(...)`"""
  
          expanded_alignments = tf.stop_gradient(tf.expand_dims(alignments, 1))
  
          history_states = tf.concat([history_states,
                                      tf.expand_dims(state, 1)], 1)
  
          # Context is the inner product of alignments and values along the
          # memory time dimension.
          # expanded_alignments shape is
          #   [batch_size, 1, memory_time]
          # history_states shape is
          #   [batch_size, memory_time, memory_size]
          # the batched matmul is over memory_time, so the output shape is
          #   [batch_size, 1, memory_size].
          # we then squeeze out the singleton dim.
  
          return tf.squeeze(tf.matmul(expanded_alignments, history_states), [1])
  
      def _prev_output(self, state, alignments, time):
          """Helper method to get previous output from memory"""
  
          # get all previous outputs from appropriate
          # attention mechanism's memory limited by current time
          prev_outputs = tf.stop_gradient(self._attention_mechanisms[
                                          self._index_of_attn_to_copy].values[
                                          :, :time, :])
  
          # multiply by alignments to get one vector from one time step
          return self._apply_alignments_to_history(alignments,
                                                   prev_outputs,
                                                   state)
  
      def _additional_likelihood(self, output, prev_output, current_time_prob):
          """Helper method to create additional likelihood to maximize"""
  
          _, likelihood = self._likelihood_fn(
              output, tf.stop_gradient(prev_output))
          return tf.where(current_time_prob < 0.5,
                          likelihood, tf.ones_like(likelihood))
  
      def _new_hidden_state(self, prev_all_cell_states,
                            new_state, alignments, time):
          """Helper method to look into rnn history"""
  
          # reshape to (batch, time, memory_time) and
          # do not include current time because
  