Returns logits for sampling ids and the next model states. Args: theta: a NestedMap of parameters. encoder_outputs: a NestedMap computed by encoder. step_ids: A tensor of shape [tgt_batch, 1]. states: A `.NestedMap` of tensors representing states that the clients
(self, theta, encoder_outputs, step_ids, states,
num_hyps_per_beam, cur_step)
source not stored for this graph (policy: none)
nothing calls this directly
no test coverage detected