| 305 | |
| 306 | |
| 307 | def pad_packed_tensor(input, lengths, value, l_min=None): |
| 308 | old_shape = input.shape |
| 309 | device = input.device |
| 310 | if not is_tensor(lengths): |
| 311 | lengths = th.tensor(lengths, dtype=th.int64, device=device) |
| 312 | else: |
| 313 | lengths = lengths.to(device) |
| 314 | max_len = as_scalar(lengths.max()) |
| 315 | |
| 316 | if l_min is not None: |
| 317 | max_len = builtins.max(max_len, l_min) |
| 318 | |
| 319 | batch_size = len(lengths) |
| 320 | x = input.new(batch_size * max_len, *old_shape[1:]) |
| 321 | x.fill_(value) |
| 322 | index = th.ones(len(input), dtype=th.int64, device=device) |
| 323 | cum_lengths = th.cumsum(lengths, 0) |
| 324 | index[cum_lengths[:-1]] += max_len - lengths[:-1] |
| 325 | index = th.cumsum(index, 0) - 1 |
| 326 | x[index] = input |
| 327 | return x.view(batch_size, max_len, *old_shape[1:]) |
| 328 | |
| 329 | |
| 330 | def pack_padded_tensor(input, lengths): |