def METHOD_NAME(
    state: train_state.TrainState,
    batch: Dict[str, Array],
    decode_rng: PRNGKey,
    ctable: CTable,
):
  """Decodes and log results for a batch."""
  inputs, outputs = batch['query'], batch['answer'][:, 1:]
  decode_rng = jax.random.fold_in(decode_rng, state.step)
  inferred = decode(state.params, inputs, decode_rng, ctable)
  questions = ctable.decode_onehot(inputs)
  infers = ctable.decode_onehot(inferred)
  goldens = ctable.decode_onehot(outputs)
  for question, inferred, golden in zip(questions, infers, goldens):
    log_decode(question, inferred, golden)