(
self,
x: torch.Tensor,
return_reg_log: bool = False,
unregularized: bool = False,
input_cp: bool = False,
output_cp: bool = False,
)
| 586 | return super().__init__(*args, **kwargs) |
| 587 | |
| 588 | def encode( |
| 589 | self, |
| 590 | x: torch.Tensor, |
| 591 | return_reg_log: bool = False, |
| 592 | unregularized: bool = False, |
| 593 | input_cp: bool = False, |
| 594 | output_cp: bool = False, |
| 595 | ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]: |
| 596 | if self.cp_size > 0 and not input_cp: |
| 597 | if not is_context_parallel_initialized: |
| 598 | initialize_context_parallel(self.cp_size) |
| 599 | |
| 600 | global_src_rank = get_context_parallel_group_rank() * self.cp_size |
| 601 | torch.distributed.broadcast(x, src=global_src_rank, group=get_context_parallel_group()) |
| 602 | |
| 603 | x = _conv_split(x, dim=2, kernel_size=1) |
| 604 | |
| 605 | if return_reg_log: |
| 606 | z, reg_log = super().encode(x, return_reg_log, unregularized) |
| 607 | else: |
| 608 | z = super().encode(x, return_reg_log, unregularized) |
| 609 | |
| 610 | if self.cp_size > 0 and not output_cp: |
| 611 | z = _conv_gather(z, dim=2, kernel_size=1) |
| 612 | |
| 613 | if return_reg_log: |
| 614 | return z, reg_log |
| 615 | return z |
| 616 | |
| 617 | def decode( |
| 618 | self, |
no test coverage detected