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

    Args:
      input_batch: The input batch, with a field `.src` and `.tgt` corresponding
        to source and target, which itself contains the `.ids` and `.paddings.`

    Returns:
      - metrics, a nestedmap of metrics.
      - results, a NestedMap like `input_batch`, with `.output_ids` (matrix of
        int ids with the decoded output) as well as the decoded length.
    """
    p = self.params
    model_theta = self.model.local_theta()
    if p.decoder.seqlen <= 0:
      raise ValueError('Must set p.decoder.seqlen > 0, current value = '
                       f'{p.decoder.seqlen}')
    batch_size = input_batch.tgt.ids.shape[0]
    decoder_state = self.model.init_states(
        inputs=input_batch.src.ids,
        input_paddings=input_batch.src.paddings,
        target_batch_size=batch_size,
        target_max_length=p.decoder.seqlen)

    global_step = base_layer.cur_global_step()

    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.model, self.model.vars_to_flax_vars(model_theta),
                         [base_layer.SCOPE_AUX_LOSS])
        new_states, xent = self.model.extend_step(states, ids)
        return new_states, xent.logits

    result = greedy_decode(
        extend_step_fn,
        decoder_state,
        input_batch.tgt.ids,
        input_batch.tgt.paddings,
        p.decoder.seqlen,
        eos_id=p.decoder.eos_id)
    # Prefix lengths are not needed for sequence model decoding.
    del result.prefix_lengths
    result.update(input_batch)
    metrics = NestedMap(
        num_decoded=(jnp.array(0.0, jnp.float32),
                     jnp.array(batch_size, jnp.float32)))
    return metrics, result