(self, f, zq)
| 388 | ) |
| 389 | |
| 390 | def forward(self, f, zq): |
| 391 | if f.shape[2] == 1 and not _USE_CP: |
| 392 | zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") |
| 393 | elif get_context_parallel_rank() == 0: |
| 394 | f_first, f_rest = f[:, :, :1], f[:, :, 1:] |
| 395 | f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] |
| 396 | zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] |
| 397 | zq_first = torch.nn.functional.interpolate(zq_first, size=f_first_size, mode="nearest") |
| 398 | zq_rest = torch.nn.functional.interpolate(zq_rest, size=f_rest_size, mode="nearest") |
| 399 | zq = torch.cat([zq_first, zq_rest], dim=2) |
| 400 | else: |
| 401 | zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") |
| 402 | |
| 403 | if self.add_conv: |
| 404 | zq = self.conv(zq) |
| 405 | |
| 406 | # f = conv_gather_from_context_parallel_region(f, dim=2, kernel_size=1) |
| 407 | norm_f = self.norm_layer(f) |
| 408 | # norm_f = conv_scatter_to_context_parallel_region(norm_f, dim=2, kernel_size=1) |
| 409 | |
| 410 | new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) |
| 411 | return new_f |
| 412 | |
| 413 | |
| 414 | def Normalize3D( |
nothing calls this directly
no test coverage detected