Split a tensor along its last dimension. Arguments: tensor: input tensor. num_partitions: number of partitions to split the tensor contiguous_split_chunks: If True, make each chunk contiguous in memory.
(tensor, num_partitions, contiguous_split_chunks=False)
| 32 | |
| 33 | |
| 34 | def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split_chunks=False): |
| 35 | """Split a tensor along its last dimension. |
| 36 | Arguments: |
| 37 | tensor: input tensor. |
| 38 | num_partitions: number of partitions to split the tensor |
| 39 | contiguous_split_chunks: If True, make each chunk contiguous |
| 40 | in memory. |
| 41 | """ |
| 42 | # Get the size and dimension. |
| 43 | last_dim = tensor.dim() - 1 |
| 44 | last_dim_size = divide(tensor.size()[last_dim], num_partitions) |
| 45 | # Split. |
| 46 | tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) |
| 47 | # Note: torch.split does not create contiguous tensors by default. |
| 48 | if contiguous_split_chunks: |
| 49 | return tuple(chunk.contiguous() for chunk in tensor_list) |
| 50 | |
| 51 | return tensor_list |
| 52 | |
| 53 | |
| 54 | class VocabUtility: |