<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def greedy_decode(logits_fn, initial_ids, temperature=0.0, initial_states=None, eos_id=EOS_ID, forced_ids=None, use_tpu=True):
 """Greedy decoding. Args: logits_fn: Interface to the model, to provide logits. Shoud take: step_num - mtf Scalar states - list of mtf.Tensor Should return: logits - [batch, vocab_size] new_states - list of mtf.Tensor temperature: a float between 0.0 (argmax) and 1.0 (random) initial_states: list of mtf.Tensor eos_id: ID for end of sentence. use_tpu: a boolean Returns: """

  length_dim = initial_ids.shape.dims[-1]
  mesh = initial_ids.mesh
  num_steps = mtf.constant(mesh, length_dim.size, dtype=tf.int32)
  def cond_fn(step_num, prev_ids, *unused_states):
    """Should we run another loop iteration."""
    overflow = mtf.equal(step_num, num_steps)
    has_eos = mtf.reduce_any(
        mtf.equal(prev_ids, eos_id), reduced_dim=length_dim)
    all_has_eos = mtf.reduce_all(has_eos)
    return mtf.logical_not(mtf.logical_or(overflow, all_has_eos))
  def body_fn(step_num, ids, *states):
    """Body function for greedy decoding.

    Args:
      step_num: a mtf.Tensor
      ids: a mtf.Tensor
      *states: additional mtf.Tensors
    Returns:
      new_step_num, new_ids, *new_states
    """
    logits, new_states = logits_fn(step_num, ids, states)
    vocab_dim = logits.shape.dims[-1]
    new_ids = mtf.sample_with_temperature(
        logits, vocab_dim, temperature)
    if forced_ids is not None:
      # force the new ids to equal the partial targets where specified
      # (positions where partial_targets contain nonzero values)
      forced = mtf.gather(forced_ids, step_num, length_dim)
      new_ids = forced + new_ids * mtf.to_int32(mtf.equal(forced, 0))
    ids += new_ids * mtf.one_hot(step_num, length_dim, dtype=tf.int32)
    new_step_num = step_num + 1
    return [new_step_num, ids] + new_states
  initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32)
  while_loop_inputs = [initial_step_num, initial_ids] + initial_states
  final_step_num, mtf_samples = mtf.while_loop(
      cond_fn, body_fn, while_loop_inputs,
      num_loop_vars=None if use_tpu else 2)[:2]
  mtf_samples = mtf.Print(mtf_samples, [final_step_num], "output_length")
  return mtf_samples