Adds relational memory to the TensorFlow graph.

    Args:
      inputs: Tensor input.
      memory: Memory output from the previous time step.
      treat_input_as_matrix: Optional, whether to treat `input` as a sequence
        of matrices. Defaulta to False, in which case the input is flattened
        into a vector.

    Returns:
      output: This time step's output.
      next_memory: The next version of memory to use.

def _build(self, inputs, memory, treat_input_as_matrix=False):
    """Adds relational memory to the TensorFlow graph.

    Args:
      inputs: Tensor input.
      memory: Memory output from the previous time step.
      treat_input_as_matrix: Optional, whether to treat `input` as a sequence
        of matrices. Defaulta to False, in which case the input is flattened
        into a vector.

    Returns:
      output: This time step's output.
      next_memory: The next version of memory to use.
    """
    if treat_input_as_matrix:
      inputs = basic.BatchFlatten(preserve_dims=2)(inputs)
      inputs_reshape = basic.BatchApply(
          basic.Linear(self._mem_size), n_dims=2)(inputs)
    else:
      inputs = basic.BatchFlatten()(inputs)
      inputs = basic.Linear(self._mem_size)(inputs)
      inputs_reshape = tf.expand_dims(inputs, 1)

    memory_plus_input = tf.concat([memory, inputs_reshape], axis=1)
    next_memory = self._attend_over_memory(memory_plus_input)

    n = inputs_reshape.get_shape().as_list()[1]
    next_memory = next_memory[:, :-n, :]

    if self._gate_style == 'unit' or self._gate_style == 'memory':
      self._input_gate, self._forget_gate = self._create_gates(
          inputs_reshape, memory)
      next_memory = self._input_gate * tf.tanh(next_memory)
      next_memory += self._forget_gate * memory

    output = basic.BatchFlatten()(next_memory)
    return output, next_memory