Initializes the sample decode state data structure. Args: inputs: An int32 tensor of shape [batch_size, length] where length <= max_decode_len. time_step: Initial time steps for decoding of shape [batch_size]. num_decodes: Number of sequences to decode per batch example.
(
*,
inputs: Tensor,
time_step: Tensor,
num_decodes: int,
max_decode_len: int,
cache: NestedTensor,
prng_key: Tensor,
pad_id: int,
token_scores: Optional[Tensor] = None,
)
| 746 | |
| 747 | |
| 748 | def _decode_init( |
| 749 | *, |
| 750 | inputs: Tensor, |
| 751 | time_step: Tensor, |
| 752 | num_decodes: int, |
| 753 | max_decode_len: int, |
| 754 | cache: NestedTensor, |
| 755 | prng_key: Tensor, |
| 756 | pad_id: int, |
| 757 | token_scores: Optional[Tensor] = None, |
| 758 | ) -> DecodingState: |
| 759 | """Initializes the sample decode state data structure. |
| 760 | |
| 761 | Args: |
| 762 | inputs: An int32 tensor of shape [batch_size, length] where length <= max_decode_len. |
| 763 | time_step: Initial time steps for decoding of shape [batch_size]. |
| 764 | num_decodes: Number of sequences to decode per batch example. |
| 765 | max_decode_len: The maximum length of the sequence to be generated (including dummy prompt |
| 766 | token). |
| 767 | cache: State of the decoder model. |
| 768 | prng_key: The initial JAX random key state. |
| 769 | pad_id: Token ID associated with padded input. |
| 770 | token_scores: Optional initial scores of shape [batch_size, length] where |
| 771 | length < max_decode_len, e.g. as produced by prefilling. Note that length should be |
| 772 | strictly less than max_decode_len, as we exclude the scores for the dummy prompt token. |
| 773 | Defaults to all zeros. |
| 774 | |
| 775 | Returns: |
| 776 | Fully initialized DecodingState. |
| 777 | |
| 778 | Raises: |
| 779 | ValueError: If inputs has an invalid shape. |
| 780 | """ |
| 781 | if inputs.shape[0] != time_step.shape[0]: |
| 782 | raise ValueError( |
| 783 | f"Expected inputs.shape[0] ({inputs.shape[0]}) " |
| 784 | f"== time_step.shape[0] ({time_step.shape[0]})." |
| 785 | ) |
| 786 | if inputs.shape[1] > max_decode_len: |
| 787 | raise ValueError( |
| 788 | f"Expected inputs.shape[1] ({inputs.shape[1]}) <= max_decode_len ({max_decode_len})." |
| 789 | ) |
| 790 | batch_size = inputs.shape[0] |
| 791 | sequences = jnp.full((batch_size, num_decodes, max_decode_len), pad_id, dtype=jnp.int32) |
| 792 | # Inputs are the prefix we will use for teacher forcing. |
| 793 | sequences = sequences.at[:, :, : inputs.shape[1]].set(inputs[:, None, :]) |
| 794 | |
| 795 | init_scores = jnp.zeros((batch_size, num_decodes, max_decode_len), dtype=jnp.float32) |
| 796 | if token_scores is not None: |
| 797 | if token_scores.shape[1] >= max_decode_len: |
| 798 | raise ValueError( |
| 799 | f"Expected token_scores.shape[1] ({token_scores.shape[1]}) < {max_decode_len}" |
| 800 | ) |
| 801 | # Note: scores at index 0 are for the dummy prompt token, which will be dropped. |
| 802 | init_scores = init_scores.at[:, :, 1 : 1 + token_scores.shape[1]].set( |
| 803 | token_scores[:, None, :] |
| 804 | ) |
| 805 |
no test coverage detected