Method ExtendStep

Computes the value vector given the query of the current step. This function is used by autoregressive decoding. Used for self-attention (hence S=T) with p.causal_masking is True. We compute the key/value/key_dists at `time_step` and cache the updated full length results in `cache_

(self,
                 theta: py_utils.NestedMap,
                 query_vec: tf.Tensor,
                 cached_states: py_utils.NestedMap,
                 paddings: Optional[tf.Tensor],
                 time_step: int,
                 segment_mask: None = None,
                 per_step_padding: None = None,
                 use_short_seq_opt: bool = False)