docs: def _build(self, x, prev_state):
      """Connects the core to the graph.
  
      Args:
        x: Input `Tensor` of shape `(batch_size, input_size)`.
        prev_state: Previous state. This could be a `Tensor`, or a tuple of
            `Tensor`s.
  
      Returns:
        The tuple `(output, state)` for this core.
  
      Raises:
        ValueError: if the `Tensor` `x` does not have rank 2.
      """
      x.get_shape().with_rank(2)
      self._batch_size = x.get_shape().as_list()[0]
      self._dtype = x.dtype
  
      x_zeros = tf.concat(
          [x, tf.zeros(
              shape=(self._batch_size, 1), dtype=self._dtype)], 1)
      x_ones = tf.concat(
          [x, tf.ones(
              shape=(self._batch_size, 1), dtype=self._dtype)], 1)
      # Weights for the halting signal
      halting_linear = basic.Linear(name="halting_linear", output_size=1)
  
      body = functools.partial(
          self._body, halting_linear=halting_linear, x_ones=x_ones)
      cumul_halting_init = tf.zeros(shape=(self._batch_size, 1),
                                    dtype=self._dtype)
      iteration_init = tf.zeros(shape=(self._batch_size, 1), dtype=self._dtype)
      core_output_size = [x.value for x in self._core.output_size]
      out_init = tf.zeros(shape=(self._batch_size,) + tuple(core_output_size),
                          dtype=self._dtype)
      cumul_state_init = _nested_zeros_like(prev_state)
      remainder_init = tf.zeros(shape=(self._batch_size, 1), dtype=self._dtype)
      (unused_final_x, final_out, unused_final_state, final_cumul_state,
       unused_final_halting, final_iteration, final_remainder) = tf.while_loop(
           self._cond, body, [x_zeros, out_init, prev_state, cumul_state_init,
                              cumul_halting_init, iteration_init, remainder_init])
  
      act_output = basic.Linear(
          name="act_output_linear", output_size=self._output_size)(final_out)
  
      return (act_output, (final_iteration, final_remainder)), final_cumul_state