| 93 | } |
| 94 | |
| 95 | static StorageView unflatten_ids(StorageView& ids, |
| 96 | const dim_t beam_size, |
| 97 | const dim_t vocabulary_size, |
| 98 | const bool is_expanded) { |
| 99 | const dim_t num_ids = ids.size(); |
| 100 | StorageView beam_origins({num_ids}, DataType::INT32); |
| 101 | |
| 102 | auto* ids_data = ids.data<int32_t>(); |
| 103 | auto* origins_data = beam_origins.data<int32_t>(); |
| 104 | |
| 105 | for (dim_t i = 0; i < num_ids; ++i) { |
| 106 | const auto flat_id = ids_data[i]; |
| 107 | const auto beam_id = flat_id / vocabulary_size; |
| 108 | const auto word_id = flat_id % vocabulary_size; |
| 109 | const auto batch_id = i / ids.dim(-1); |
| 110 | ids_data[i] = word_id; |
| 111 | origins_data[i] = is_expanded ? batch_id * beam_size + beam_id : batch_id; |
| 112 | } |
| 113 | |
| 114 | return beam_origins; |
| 115 | } |
| 116 | |
| 117 | static void append_step_output(StorageView& history, // [batch, beam, time, ...] |
| 118 | StorageView step_output, // [batch, beam, ...] |