code: def _build(self, inputs, memory, treat_input_as_matrix=False):
      
      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
docs: 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.