(params, indices)
| 139 | |
| 140 | |
| 141 | def gather_nd(params, indices): |
| 142 | indices = indices.t().long() |
| 143 | ndim = indices.size(0) |
| 144 | idx = torch.zeros_like(indices[0]).long() |
| 145 | m = 1 |
| 146 | |
| 147 | for i in range(ndim)[::-1]: |
| 148 | idx += indices[i] * m |
| 149 | m *= params.size(i) |
| 150 | |
| 151 | params = params.reshape((-1, *tuple(torch.tensor(params.size()[ndim:])))) |
| 152 | return params[idx] |
| 153 | |
| 154 | |
| 155 | def gather_tree(step_ids, parent_ids, max_sequence_lengths, end_token): |