(
self,
x: torch.Tensor,
return_reg_log: bool = False,
unregularized: bool = False,
input_cp: bool = False,
output_cp: bool = False,
use_cp: bool = True,
)
| 560 | return super().__init__(*args, **kwargs) |
| 561 | |
| 562 | def encode( |
| 563 | self, |
| 564 | x: torch.Tensor, |
| 565 | return_reg_log: bool = False, |
| 566 | unregularized: bool = False, |
| 567 | input_cp: bool = False, |
| 568 | output_cp: bool = False, |
| 569 | use_cp: bool = True, |
| 570 | ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]: |
| 571 | if self.cp_size <= 1: |
| 572 | use_cp = False |
| 573 | if self.cp_size > 0 and use_cp and not input_cp: |
| 574 | if not is_context_parallel_initialized: |
| 575 | initialize_context_parallel(self.cp_size) |
| 576 | |
| 577 | global_src_rank = get_context_parallel_group_rank() * self.cp_size |
| 578 | torch.distributed.broadcast(x, src=global_src_rank, group=get_context_parallel_group()) |
| 579 | |
| 580 | x = _conv_split(x, dim=2, kernel_size=1) |
| 581 | |
| 582 | if return_reg_log: |
| 583 | z, reg_log = super().encode(x, return_reg_log, unregularized, use_cp=use_cp) |
| 584 | else: |
| 585 | z = super().encode(x, return_reg_log, unregularized, use_cp=use_cp) |
| 586 | |
| 587 | if self.cp_size > 0 and use_cp and not output_cp: |
| 588 | z = _conv_gather(z, dim=2, kernel_size=1) |
| 589 | |
| 590 | if return_reg_log: |
| 591 | return z, reg_log |
| 592 | return z |
| 593 | |
| 594 | def decode( |
| 595 | self, |
no test coverage detected