//  [python]
def build_sample_graph(self, input_pianorolls=None, outer_masks=None,
                         total_gibbs_steps=None):
    """Builds the tf.while_loop based sampling graph.

    Args:
      input_pianorolls: Optional input pianorolls override. If None, uses the
          pianorolls placeholder.
      outer_masks: Optional input outer_masks override. If None, uses the
          outer_masks placeholder.
      total_gibbs_steps: Optional input total_gibbs_steps override. If None,
          uses the total_gibbs_steps placeholder.
    Returns:
      The output op of the graph.
    """
    if input_pianorolls is None:
      input_pianorolls = self.inputs["pianorolls"]
    if outer_masks is None:
      outer_masks = self.inputs["outer_masks"]

    tt = tf.shape(input_pianorolls)[1]
    sample_steps = tf.cast(self.inputs["sample_steps"], tf.float32)
    if total_gibbs_steps is None:
      total_gibbs_steps = self.inputs["total_gibbs_steps"]
    temperature = self.inputs["temperature"]

    input_pianorolls = tf.cast(input_pianorolls, tf.float32)
    outer_masks = self.make_outer_masks(outer_masks, input_pianorolls)

    # Calculate total_gibbs_steps as steps * num_instruments if not given.
    total_gibbs_steps = tf.cond(
        tf.equal(total_gibbs_steps, 0),
        lambda: tf.cast(tt * self.hparams.num_instruments, tf.float32),
        lambda: tf.cast(total_gibbs_steps, tf.float32))

    # sample_steps is set to total_gibbs_steps if not given.
    sample_steps = tf.cond(
        tf.equal(sample_steps, 0),
        lambda: total_gibbs_steps,
        lambda: tf.cast(sample_steps, tf.float32))

    def infer_step(pianorolls, step_count):
      """Called by tf.while_loop, takes a Gibbs step."""
      mask_prob = compute_mask_prob_from_yao_schedule(step_count,
                                                      total_gibbs_steps)
      # 1 indicates mask out, 0 is not mask.
      masks = make_bernoulli_masks(tf.shape(pianorolls), mask_prob,
                                   outer_masks)

      logits = self.predict(pianorolls, masks)
      samples = sample_with_temperature(logits, temperature=temperature)

      outputs = pianorolls * (1 - masks) + samples * masks

      check_completion_op = tf.assert_equal(
          tf.where(tf.equal(tf.reduce_max(masks, axis=2), 1.),
                   tf.reduce_max(outputs, axis=2),
                   tf.reduce_max(pianorolls, axis=2)),
          1.)
      with tf.control_dependencies([check_completion_op]):
        outputs = tf.identity(outputs)

      step_count += 1
      return outputs, step_count

    current_step = tf.cast(self.inputs["current_step"], tf.float32)

    # Initializes pianorolls by evaluating the model once to fill in all gaps.
    logits = self.predict(tf.cast(input_pianorolls, tf.float32), outer_masks)
    samples = sample_with_temperature(logits, temperature=temperature)
    tf.get_variable_scope().reuse_variables()

    self.samples, current_step = tf.while_loop(
        lambda samples, current_step: current_step < sample_steps,
        infer_step, [samples, current_step],
        shape_invariants=[
            tf.TensorShape([None, None, None, None]),
            tf.TensorShape(None),
        ],
        back_prop=False,
        parallel_iterations=1,
        name="coco_while")
    self.samples.set_shape(input_pianorolls.shape)
    return self.samples