Method ExtendStep

Transformer decoder layer, extend one step in autoregressive decoding. query_vec and aux_* may have different batch sizes, e.g., during a beam search. target_batch must be a multiple of source_batch and query_vec[i * batch_multiplier + j] corresponds to aux_vec[i], where batch_multi

(
      self,
      theta,
      query_vec,
      aux_vec,
      aux_paddings,
      cached_states,
      time_step,
      use_short_seq_opt=False,
      per_step_padding=None,
      *,
      compute_atten_probs=False,
      segment_mask=None,
      aux_segment_mask=None,
  )