(
self,
z: torch.Tensor,
input_cp: bool = False,
output_cp: bool = False,
split_kernel_size: int = 1,
**kwargs,
)
| 615 | return z |
| 616 | |
| 617 | def decode( |
| 618 | self, |
| 619 | z: torch.Tensor, |
| 620 | input_cp: bool = False, |
| 621 | output_cp: bool = False, |
| 622 | split_kernel_size: int = 1, |
| 623 | **kwargs, |
| 624 | ): |
| 625 | if self.cp_size > 0 and not input_cp: |
| 626 | if not is_context_parallel_initialized: |
| 627 | initialize_context_parallel(self.cp_size) |
| 628 | |
| 629 | global_src_rank = get_context_parallel_group_rank() * self.cp_size |
| 630 | torch.distributed.broadcast(z, src=global_src_rank, group=get_context_parallel_group()) |
| 631 | |
| 632 | z = _conv_split(z, dim=2, kernel_size=split_kernel_size) |
| 633 | |
| 634 | x = super().decode(z, **kwargs) |
| 635 | |
| 636 | if self.cp_size > 0 and not output_cp: |
| 637 | x = _conv_gather(x, dim=2, kernel_size=split_kernel_size) |
| 638 | |
| 639 | return x |
| 640 | |
| 641 | def forward( |
| 642 | self, |
no test coverage detected