Holds sample decoding state data.
| 729 | |
| 730 | |
| 731 | class DecodingState(NamedTuple): |
| 732 | """Holds sample decoding state data.""" |
| 733 | |
| 734 | # The position of the decoding loop in the length dimension. |
| 735 | cur_index: Tensor # scalar int32: current decoded length index. |
| 736 | # The active sequences. |
| 737 | sequences: Tensor # int32: [batch_size, num_decodes, max_decode_len]. |
| 738 | # The sequence token log probabilities. |
| 739 | token_scores: Tensor # float32: [batch_size, num_decodes, max_decode_len]. |
| 740 | # Whether a stop decoding condition has been reached. |
| 741 | stop_decoding: Tensor # bool: [batch_size, num_decodes]. |
| 742 | # The current state of the autoregressive decoding caches. |
| 743 | cache: NestedTensor |
| 744 | # Random generator state. |
| 745 | prng_key: Tensor |
| 746 | |
| 747 | |
| 748 | def _decode_init( |
no outgoing calls
no test coverage detected