code: def decode_step(self, step_input, states):
          
          step_output, states, step_additional_outputs =\
              self.decoder(self.tgt_embed(step_input), states)
          step_output = self.tgt_proj(step_output)
          return step_output, states, step_additional_outputs
docs: One step decoding of the translation model.
  
          Parameters
          ----------
          step_input : NDArray
              Shape (batch_size,)
          states : list of NDArrays
  
          Returns
          -------
          step_output : NDArray
              Shape (batch_size, C_out)
          states : list
          step_additional_outputs : list
              Additional outputs of the step, e.g, the attention weights