Output values after performing sample decoding.
| 815 | |
| 816 | |
| 817 | class SampleOutputs(flax_struct.PyTreeNode): |
| 818 | """Output values after performing sample decoding.""" |
| 819 | |
| 820 | # Sequences which may or may not end with eos_id. |
| 821 | sequences: Tensor # int32: [batch_size, num_decodes, max_decode_len]. |
| 822 | # Scores corresponding to sequences above. |
| 823 | token_scores: Tensor # float32: [batch_size, num_decodes, max_decode_len]. |
| 824 | |
| 825 | |
| 826 | class StopDecodingCondition(Protocol): |