(input_, dim)
| 74 | |
| 75 | |
| 76 | def _split(input_, dim): |
| 77 | cp_world_size = get_context_parallel_world_size() |
| 78 | |
| 79 | if cp_world_size == 1: |
| 80 | return input_ |
| 81 | |
| 82 | cp_rank = get_context_parallel_rank() |
| 83 | |
| 84 | # print('in _split, cp_rank:', cp_rank, 'input_size:', input_.shape) |
| 85 | |
| 86 | inpu_first_frame_ = input_.transpose(0, dim)[:1].transpose(0, dim).contiguous() |
| 87 | input_ = input_.transpose(0, dim)[1:].transpose(0, dim).contiguous() |
| 88 | dim_size = input_.size()[dim] // cp_world_size |
| 89 | |
| 90 | input_list = torch.split(input_, dim_size, dim=dim) |
| 91 | output = input_list[cp_rank] |
| 92 | |
| 93 | if cp_rank == 0: |
| 94 | output = torch.cat([inpu_first_frame_, output], dim=dim) |
| 95 | output = output.contiguous() |
| 96 | |
| 97 | # print('out _split, cp_rank:', cp_rank, 'output_size:', output.shape) |
| 98 | |
| 99 | return output |
| 100 | |
| 101 | |
| 102 | def _gather(input_, dim): |
nothing calls this directly
no test coverage detected