Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or nested list/tuples/dict of tensors.
(tensors, new_tensors, padding_index=-100)
| 62 | |
| 63 | |
| 64 | def nested_concat(tensors, new_tensors, padding_index=-100): |
| 65 | """ |
| 66 | Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or |
| 67 | nested list/tuples/dict of tensors. |
| 68 | """ |
| 69 | assert type(tensors) == type( |
| 70 | new_tensors |
| 71 | ), f"Expected `tensors` and `new_tensors` to have the same type but found {type(tensors)} and {type(new_tensors)}." |
| 72 | if isinstance(tensors, (list, tuple)): |
| 73 | return type(tensors)(nested_concat(t, n, padding_index=padding_index) for t, n in zip(tensors, new_tensors)) |
| 74 | elif isinstance(tensors, dict): |
| 75 | assert set(tensors.keys()) == set(new_tensors.keys()) |
| 76 | return type(tensors)({k: nested_concat(tensors[k], new_tensors[k], padding_index=padding_index) for k in tensors.keys()}) |
| 77 | elif isinstance(tensors, torch.Tensor): |
| 78 | return torch_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index) |
| 79 | elif isinstance(tensors, np.ndarray): |
| 80 | return numpy_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index) |
| 81 | elif tensors is None: |
| 82 | return None |
| 83 | else: |
| 84 | raise TypeError(f"Unsupported type for concatenation: got {type(tensors)}") |
| 85 | |
| 86 | |
| 87 | def nested_cpu(tensors): |