def decode(self, input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
    """Greedy decodes the input_batch.

    Args:
      input_batch: The input batch, with fields like `.ids`.

    Returns:
      - metrics, a NestedMap containing str keys and (metrics, weight) pairs.
      - A NestedMap like `input_batch`, with `.prefix_lengths` (vector of
        randomly generated ints indicating the lengths of prefixes for each
        row), and `.output_ids` (matrix of int ids with the decoded output).
    """
    p = self.params
    if p.decoder.seqlen <= 0:
      raise ValueError('Must set p.decoder.seqlen > 0, current value = '
                       f'{p.decoder.seqlen}')
    batch_size = input_batch.ids.shape[0]
    maxval = jnp.sum(1 - input_batch.paddings, axis=1).astype(jnp.int32)
    minval = jnp.minimum(maxval, p.decoder.min_prefix_len)
    prefix_lengths = jax.random.randint(base_layer.next_prng_key(),
                                        [batch_size], minval, maxval + 1,
                                        input_batch.ids.dtype)
    decoder_state = self.lm.init_states(
        target_batch_size=batch_size,
        target_max_length=p.decoder.seqlen)

    global_step = base_layer.cur_global_step()

    lm_theta = self.lm.local_theta()
    def extend_step_fn(states, ids):
      with base_layer.JaxContext.new_context(
          prng_key=base_layer.next_prng_key(),
          global_step=global_step) as jax_context:
        jax_context.bind(self.lm, self.lm.vars_to_flax_vars(lm_theta),
                         [base_layer.SCOPE_AUX_LOSS])
        new_states, xent = self.lm.extend_step(states, ids)
        return new_states, xent.logits

    result = greedy_decode(
        extend_step_fn,
        decoder_state,
        input_batch.ids,
        input_batch.paddings,
        p.decoder.seqlen,
        max_decode_steps=p.decoder.max_decode_steps,
        prefix_lengths=prefix_lengths,
        eos_id=p.decoder.eos_id)
    result.update(input_batch)

    metrics = NestedMap(
        num_decoded=(jnp.array(0.0, jnp.float32),
                     jnp.array(batch_size, jnp.float32)))
    return metrics, result