See `BaseDecoder.forward_step` for details.
(
self,
*,
cached_states: Nested[Tensor],
input_batch: Nested[Tensor],
is_prefill: bool = False,
**kwargs,
)
| 766 | return states, outputs |
| 767 | |
| 768 | def extend_step( |
| 769 | self, |
| 770 | *, |
| 771 | cached_states: Nested[Tensor], |
| 772 | input_batch: Nested[Tensor], |
| 773 | is_prefill: bool = False, |
| 774 | **kwargs, |
| 775 | ) -> tuple[Nested[Tensor], Nested[Tensor]]: |
| 776 | """See `BaseDecoder.forward_step` for details.""" |
| 777 | cfg = self.config |
| 778 | mode = ForwardMode.PREFILL if is_prefill else ForwardMode.EXTEND_STEP |
| 779 | time_step: Tensor = cached_states["time_step"] |
| 780 | assert time_step.ndim == 1 |
| 781 | |
| 782 | input_ids: Tensor = input_batch["input_ids"] |
| 783 | batch, step = input_ids.shape |
| 784 | # This position is for embedding, not attention. |
| 785 | positions = repeat(jnp.arange(step), "t -> b t", b=batch) + time_step[:, None] # [B, T] |
| 786 | |
| 787 | # TODO(dhwang2): self_attention_logit_biases is used by only T5. Delete all this |
| 788 | # self_attention_logit_biases mess. |
| 789 | if cfg.attention_mask is not None: |
| 790 | cached_inputs: Tensor = cached_states["input_ids"] |
| 791 | target_len = cached_inputs.shape[-1] |
| 792 | # [B, step, T] |
| 793 | oh_indices = jax.nn.one_hot(positions, target_len, dtype=input_ids.dtype) |
| 794 | keep_mask = ~oh_indices.any(axis=1) # [B, T] |
| 795 | input_ids_scattered = jnp.einsum("bs,bst->bt", input_ids, oh_indices) |
| 796 | updated_inputs = cached_inputs * keep_mask + input_ids_scattered |
| 797 | |
| 798 | # Compute self-attention-mask logit biases. [B, N, T, T]. |
| 799 | self_attention_biases = self.compute_attention_logit_biases( |
| 800 | updated_inputs, |
| 801 | segment_ids=jnp.ones_like(updated_inputs), |
| 802 | positions=jnp.arange(target_len)[None, :], |
| 803 | ) |
| 804 | # Select logit biases corresponding to time step. [B, N, step, T]. |
| 805 | if self_attention_biases is not None: |
| 806 | self_attention_biases = jnp.take_along_axis( |
| 807 | self_attention_biases, |
| 808 | positions[:, None, :, None], |
| 809 | axis=2, |
| 810 | mode="clip", |
| 811 | ) |
| 812 | else: |
| 813 | self_attention_biases = None |
| 814 | |
| 815 | if "positions" in kwargs: |
| 816 | raise ValueError("positions is supported only in FORWARD.") |
| 817 | |
| 818 | input_segment_ids = input_batch.get("input_segment_ids", None) |
| 819 | updated_states, outputs = self._forward_for_mode( |
| 820 | mode=mode, |
| 821 | input_batch={**input_batch, "positions": positions}, # emb may use positional encoding |
| 822 | self_attention_logit_biases=self_attention_biases, |
| 823 | cached_states=cached_states, |
| 824 | **kwargs, |
| 825 | ) |
no test coverage detected