Convert a list of 1d tensors into a padded 2d tensor.
(values, pad_idx=0, left_pad=False, shift_right=False, max_len=None, shift_id=1)
| 42 | |
| 43 | |
| 44 | def collate_1d(values, pad_idx=0, left_pad=False, shift_right=False, max_len=None, shift_id=1): |
| 45 | """Convert a list of 1d tensors into a padded 2d tensor.""" |
| 46 | size = max(v.size(0) for v in values) if max_len is None else max_len |
| 47 | res = values[0].new(len(values), size).fill_(pad_idx) |
| 48 | |
| 49 | def copy_tensor(src, dst): |
| 50 | assert dst.numel() == src.numel() |
| 51 | if shift_right: |
| 52 | dst[1:] = src[:-1] |
| 53 | dst[0] = shift_id |
| 54 | else: |
| 55 | dst.copy_(src) |
| 56 | |
| 57 | for i, v in enumerate(values): |
| 58 | copy_tensor(v, res[i][size - len(v):] if left_pad else res[i][:len(v)]) |
| 59 | return res |
| 60 | |
| 61 | |
| 62 | def collate_2d(values, pad_idx=0, left_pad=False, shift_right=False, max_len=None): |
nothing calls this directly
no test coverage detected