code: 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
              ids - mtf Tensor with shape [..., length]
              states - list of mtf.Tensor
            Should return:
              logits - [batch, vocab_size]
              new_states - list of mtf.Tensor
        initial_ids: mtf.Tensor with shape [..., length], containing zeros.
        temperature: a float between 0.0 (argmax) and 1.0 (random)
        initial_states: list of mtf.Tensor
        eos_id: ID for end of sentence.
        forced_ids: optional mtf.Tensor with shape [..., length]
        use_tpu: a boolean
      Returns:
        Tensor with shape [..., length]
      """
      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
  
file_path: mesh/mesh_tensorflow/beam_search.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  from mesh_tensorflow import ops_with_redefined_builtins
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: greedy_decode(logits_fn, initial_ids, temperature, initial_states, eos_id, forced_ids, use_tpu) and returns: mtf_samples, mtf, unknown