<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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