def predict_step(batch, state, cache, eos_idx, config):
  """Compute the prediction for the given model in inference mode."""

  logging.info('predict_step(batch=%s)', batch)
  variables = {'params': state.optimizer.target}
  model = models.Model(config)
  encoded, encoded_mask = model.apply(
      variables, batch, method=models.Model.encode)

  encoded_inputs = decode.flat_batch_beam_expand(encoded, config.beam_size)
  encoded_mask = decode.flat_batch_beam_expand(encoded_mask, config.beam_size)

  def tokens_ids_to_logits(flat_ids, flat_cache):
    """Token slice to logits from decoder model."""
    # --> [batch * beam, 1, vocab]
    flat_logits, new_vars = model.apply(
        {
            'params': state.optimizer.target,
            'cache': flat_cache
        },
        flat_ids,
        encoded_inputs,
        flat_ids > 0,
        encoded_mask,
        mutable=['cache'],
        method=models.Model.decode)
    new_flat_cache = new_vars['cache']
    # Remove singleton sequence-length dimension:
    # [batch * beam, 1, vocab] --> [batch * beam, vocab]
    flat_logits = flat_logits.squeeze(axis=1)
    return flat_logits, new_flat_cache

  # Using the above-defined single-step decoder function, run a
  # beam search over possible sequences given input encoding.
  beam_seqs, _, = decode.beam_search(
      batch['token'],
      cache,
      tokens_ids_to_logits,
      beam_size=config.beam_size,
      alpha=0.6,
      eos_id=eos_idx,
      max_decode_len=config.max_decode_step)
  # Beam search returns [n_batch, n_beam, n_length + 1] with beam dimension
  # sorted in increasing order of log-probability.
  # Return the highest scoring beam sequence.
  return beam_seqs[:, -1]
Compute the prediction for the given model in inference mode.
language/gscan/xattn_model/predict.py