jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_upsample_nearest2d(self):
        x = torch.randn(2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test channels last
        x2 = x.clone().detach().to(memory_format=torch.channels_last).requires_grad_()
        y2 = t.interpolate(x2, scale_factor=2, mode="nearest")
        y2.mean().backward()
        self.assertTrue(y2.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(y1, y2)
        self.assertTrue(x2.grad.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(x1.grad, x2.grad)

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last)
                    .requires_grad_()
                )
                y5 = t.interpolate(x5, scale_factor=2, mode="nearest")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last)
                )
//  [python]
def test_upsample_nearest2d(self):
        x = torch.randn(2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test channels last
        x2 = x.clone().detach().to(memory_format=torch.channels_last).requires_grad_()
        y2 = F.interpolate(x2, scale_factor=2, mode="nearest")
        y2.mean().backward()
        self.assertTrue(y2.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(y1, y2)
        self.assertTrue(x2.grad.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(x1.grad, x2.grad)

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last)
                    .requires_grad_()
                )
                y5 = F.interpolate(x5, scale_factor=2, mode="nearest")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last)
                )
//  [python]
def test_upsample_nearest3d(self):
        x = torch.randn(2, 2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last_3d)
                    .requires_grad_()
                )
                y5 = t.interpolate(x5, scale_factor=2, mode="nearest")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last_3d))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last_3d)
                )
//  [python]
def test_upsample_nearest3d(self):
        x = torch.randn(2, 2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last_3d)
                    .requires_grad_()
                )
                y5 = F.interpolate(x5, scale_factor=2, mode="nearest")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last_3d))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last_3d)
                )
//  [python]
def grid_sample(self, pos: torch.Tensor):
        # pos: (B, 2, H, W)
        B, _, H, W = pos.size()
        # grid = make_coords_grid(B, H//10, W//10, pos.device)
        # grid = grid / grid.max(dim=1, keepdim=True).values * 2 - 1
        # print("grid min = {}, grid max = {}".format(grid.min().item(), grid.max().item()))
        # grid = rearrange(grid, "b c h w -> b h w c")
        # sampled_grid = F.grid_sample(pos, grid, align_corners=False)
        sampled_grid = F.upsample(pos, scale_factor=0.1, mode="bilinear")
        # sampled_grid = F.interpolate(pos, scale_factor=0.1, mode="area")
        # print("sample_grid size = {}".format(sampled_grid.size()))
        return sampled_grid
//  [python]
def grid_sample(self, pos: torch.Tensor):
        # pos: (B, 2, H, W)
        B, _, H, W = pos.size()
        # grid = make_coords_grid(B, H//10, W//10, pos.device)
        # grid = grid / grid.max(dim=1, keepdim=True).values * 2 - 1
        # print("grid min = {}, grid max = {}".format(grid.min().item(), grid.max().item()))
        # grid = rearrange(grid, "b c h w -> b h w c")
        # sampled_grid = F.grid_sample(pos, grid, align_corners=False)
        sampled_grid = F.interpolate(pos, scale_factor=0.1, mode="bilinear", align_corners=False)
        # sampled_grid = F.interpolate(pos, scale_factor=0.1, mode="area")
        # print("sample_grid size = {}".format(sampled_grid.size()))
        return sampled_grid
//  [python]
def test_upsample_trilinear3d(self):
        x = torch.randn(2, 2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="trilinear")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="trilinear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.02)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="trilinear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last_3d)
                    .requires_grad_()
                )
                y5 = t.interpolate(x5, scale_factor=2, mode="trilinear")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last_3d))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last_3d)
                )
//  [python]
def test_upsample_trilinear3d(self):
        x = torch.randn(2, 2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="trilinear")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="trilinear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.02)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="trilinear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last_3d)
                    .requires_grad_()
                )
                y5 = F.interpolate(x5, scale_factor=2, mode="trilinear")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last_3d))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last_3d)
                )
//  [python]
def forward(self, imgs):
        imgs = F.upsample(
            imgs,
            size=(self.input_resolution, ) * 2,
            mode='bilinear')
        return super().forward(imgs).view(imgs.shape[0], -1)
//  [python]
def forward(self, imgs):
        imgs = F.interpolate(
            imgs,
            size=(self.input_resolution, ) * 2,
            mode='bilinear',
            align_corners=False)
        return super().forward(imgs).view(imgs.shape[0], -1)
//  [python]
def test_upsample_bilinear2d(self):
        x = torch.randn(2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="bilinear")
        y1.mean().backward()

        # test channels last
        x2 = x.clone().detach().to(memory_format=torch.channels_last).requires_grad_()
        y2 = t.interpolate(x2, scale_factor=2, mode="bilinear")
        y2.mean().backward()
        self.assertTrue(y2.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(y1, y2)
        self.assertTrue(x2.grad.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(x1.grad, x2.grad)

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="bilinear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="bilinear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last)
                    .requires_grad_()
                )
                y5 = t.interpolate(x5, scale_factor=2, mode="bilinear")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last)
                )
//  [python]
def test_upsample_bilinear2d(self):
        x = torch.randn(2, 2, 4, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="bilinear")
        y1.mean().backward()

        # test channels last
        x2 = x.clone().detach().to(memory_format=torch.channels_last).requires_grad_()
        y2 = F.interpolate(x2, scale_factor=2, mode="bilinear")
        y2.mean().backward()
        self.assertTrue(y2.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(y1, y2)
        self.assertTrue(x2.grad.is_contiguous(memory_format=torch.channels_last))
        self.assertEqual(x1.grad, x2.grad)

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="bilinear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="bilinear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)

                x5 = (
                    x.clone()
                    .detach()
                    .to(datatype)
                    .to(memory_format=torch.channels_last)
                    .requires_grad_()
                )
                y5 = F.interpolate(x5, scale_factor=2, mode="bilinear")
                y5.mean().backward()
                self.assertTrue(y5.dtype == datatype)
                self.assertTrue(x5.grad.dtype == datatype)
                self.assertTrue(y5.is_contiguous(memory_format=torch.channels_last))
                self.assertTrue(
                    x5.grad.is_contiguous(memory_format=torch.channels_last)
                )
//  [python]
def preprocess_generated_image(image):
    transform_list = transforms.Normalize(NORM_MEAN, NORM_STD)
    image = transform_list(image * 0.5 + 0.5)
    image = t.nn.functional.interpolate(
        image, 224, mode="bicubic", align_corners=True
    )
    return image
//  [python]
def preprocess_generated_image(image):
    transform_list = transforms.Normalize(NORM_MEAN, NORM_STD)
    image = transform_list(image * 0.5 + 0.5)
    image = torch.nn.functional.interpolate(
        image, 224, mode="bicubic", align_corners=True
    )
    return image
//  [python]
def preprocess_input_image(input_image_path, size):
    pil_image = Image_PIL.open(input_image_path).convert("RGB")
    transform_list = transforms.Compose(
        [
            data_utils.CenterCropLongEdge(),
            transforms.Resize((size, size)),
            transforms.ToTensor(),
            transforms.Normalize(NORM_MEAN, NORM_STD),
        ]
    )
    tensor_image = transform_list(pil_image)
    tensor_image = t.nn.functional.interpolate(
        tensor_image.unsqueeze(0), 224, mode="bicubic", align_corners=True
    )
    return tensor_image
//  [python]
def preprocess_input_image(input_image_path, size):
    pil_image = Image_PIL.open(input_image_path).convert("RGB")
    transform_list = transforms.Compose(
        [
            data_utils.CenterCropLongEdge(),
            transforms.Resize((size, size)),
            transforms.ToTensor(),
            transforms.Normalize(NORM_MEAN, NORM_STD),
        ]
    )
    tensor_image = transform_list(pil_image)
    tensor_image = torch.nn.functional.interpolate(
        tensor_image.unsqueeze(0), 224, mode="bicubic", align_corners=True
    )
    return tensor_image
//  [python]
def __call__(self, x_t, attention_store):
        k = 1
        maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
        maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, self.max_num_words) for item in maps]
        maps = torch.cat(maps, dim=1)
        maps = (maps * self.alpha_layers).sum(-1).mean(1)
        mask = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
        mask = nnf.upsample(mask, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.threshold)
        mask = (mask[:1] + mask[1:]).float()
        x_t = x_t[:1] + mask * (x_t - x_t[:1])
        return x_t
//  [python]
def __call__(self, x_t, attention_store):
        k = 1
        maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
        maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, self.max_num_words) for item in maps]
        maps = torch.cat(maps, dim=1)
        maps = (maps * self.alpha_layers).sum(-1).mean(1)
        mask = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
        mask = nnf.interpolate(mask, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.threshold)
        mask = (mask[:1] + mask[1:]).float()
        x_t = x_t[:1] + mask * (x_t - x_t[:1])
        return x_t
//  [python]
def forward(self, inputs):
        """Forward function.

        Args:
            inputs (dict): Output dict of FbaEncoder.
        Returns:
            tuple(Tensor): Predicted alpha, fg and bg of the current batch.
        """

        conv_out = inputs['conv_out']
        img = inputs['merged']
        two_channel_trimap = inputs['two_channel_trimap']
        conv5 = conv_out[-1]
        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(
                t.interpolate(
                    pool_scale(conv5), (input_size[2], input_size[3]),
                    mode='bilinear',
                    align_corners=self.align_corners))
        ppm_out = torch.cat(ppm_out, 1)
        x = self.conv_up1(ppm_out)

        x = t.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-4]), 1)

        x = self.conv_up2(x)
        x = t.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-5]), 1)
        x = self.conv_up3(x)

        x = t.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-6][:, :3], img, two_channel_trimap), 1)
        output = self.conv_up4(x)
        alpha = torch.clamp(output[:, 0:1], 0, 1)
        F = torch.sigmoid(output[:, 1:4])
        B = torch.sigmoid(output[:, 4:7])

        return alpha, F, B
//  [python]
def forward(self, inputs):
        """Forward function.

        Args:
            inputs (dict): Output dict of FbaEncoder.
        Returns:
            tuple(Tensor): Predicted alpha, fg and bg of the current batch.
        """

        conv_out = inputs['conv_out']
        img = inputs['merged']
        two_channel_trimap = inputs['two_channel_trimap']
        conv5 = conv_out[-1]
        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(
                nn.functional.interpolate(
                    pool_scale(conv5), (input_size[2], input_size[3]),
                    mode='bilinear',
                    align_corners=self.align_corners))
        ppm_out = torch.cat(ppm_out, 1)
        x = self.conv_up1(ppm_out)

        x = torch.nn.functional.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-4]), 1)

        x = self.conv_up2(x)
        x = torch.nn.functional.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-5]), 1)
        x = self.conv_up3(x)

        x = torch.nn.functional.interpolate(
            x,
            scale_factor=2,
            mode='bilinear',
            align_corners=self.align_corners)

        x = torch.cat((x, conv_out[-6][:, :3], img, two_channel_trimap), 1)
        output = self.conv_up4(x)
        alpha = torch.clamp(output[:, 0:1], 0, 1)
        F = torch.sigmoid(output[:, 1:4])
        B = torch.sigmoid(output[:, 4:7])

        return alpha, F, B
//  [python]
def render_validation_views(*,
                            render_size,
                            max_size,
                            thetas=(30,),
                            phis=(-25,),
                            device,
                            depth_cmap="jet",
                            **render_kwargs):
  """Render a frame from a camera at a new perspective."""
  height, width, focal = scene.scale_intrinsics(min(render_size, max_size))

  frames = []
  for theta in thetas:
    for phi in phis:
      cam2world = scene.pose_spherical(theta=theta, phi=phi, radius=4)
      rays = scene.camera_rays(cam2world, height, width, focal)
      rays = [torch.from_numpy(r).to(device).type(torch.float32) for r in rays]
      rendered, _ = render_rays(
          rays, white_bkgd=True, device=device,
          **render_kwargs)  # rgb, depth, disparity, silhouette.
      rendered = torch.cat(rendered, dim=-1)  # [H, W, 6].
      frames.append(rendered)

  frames = torch.stack(frames)
  if height != render_size:
    frames = frames.movedim(-1, 1)  # NHWC to NCHW.
    frames = t.nn.functional.interpolate(frames, render_size, mode="bilinear")
    frames = frames.movedim(1, -1)  # NHWC.

  frames = frames.cpu()
  rgb, depth, disparity, silhouette = torch.split(frames, [3, 1, 1, 1], dim=-1)
  depth_cmap = plt.get_cmap(depth_cmap)
  return (
      rgb.numpy(),  # [T, H, W, 3].
      depth_cmap(depth.numpy().squeeze(-1))[Ellipsis, :3],  # [T, H, W, 1].
      depth_cmap(disparity.numpy().squeeze(-1))[Ellipsis, :3],  # [T, H, W, 1].
      silhouette.numpy())
//  [python]
def render_validation_views(*,
                            render_size,
                            max_size,
                            thetas=(30,),
                            phis=(-25,),
                            device,
                            depth_cmap="jet",
                            **render_kwargs):
  """Render a frame from a camera at a new perspective."""
  height, width, focal = scene.scale_intrinsics(min(render_size, max_size))

  frames = []
  for theta in thetas:
    for phi in phis:
      cam2world = scene.pose_spherical(theta=theta, phi=phi, radius=4)
      rays = scene.camera_rays(cam2world, height, width, focal)
      rays = [torch.from_numpy(r).to(device).type(torch.float32) for r in rays]
      rendered, _ = render_rays(
          rays, white_bkgd=True, device=device,
          **render_kwargs)  # rgb, depth, disparity, silhouette.
      rendered = torch.cat(rendered, dim=-1)  # [H, W, 6].
      frames.append(rendered)

  frames = torch.stack(frames)
  if height != render_size:
    frames = frames.movedim(-1, 1)  # NHWC to NCHW.
    frames = F.interpolate(frames, render_size, mode="bilinear")
    frames = frames.movedim(1, -1)  # NHWC.

  frames = frames.cpu()
  rgb, depth, disparity, silhouette = torch.split(frames, [3, 1, 1, 1], dim=-1)
  depth_cmap = plt.get_cmap(depth_cmap)
  return (
      rgb.numpy(),  # [T, H, W, 3].
      depth_cmap(depth.numpy().squeeze(-1))[Ellipsis, :3],  # [T, H, W, 1].
      depth_cmap(disparity.numpy().squeeze(-1))[Ellipsis, :3],  # [T, H, W, 1].
      silhouette.numpy())
//  [python]
def resize(x, scale_factor):
    return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False)
//  [python]
def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.lrelu(self.conv_up2(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H*self.upscale, :W*self.upscale]
//  [python]
def forward(self, x):
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H*self.upscale, :W*self.upscale]
//  [python]
def forward(self, x, x_skip):
        x = t.interpolate(x, scale_factor=2, mode="trilinear", align_corners=True)
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv_block(x)
        return x
//  [python]
def forward(self, x, x_skip):
        x = nn.functional.interpolate(x, scale_factor=2, mode="trilinear", align_corners=True)
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv_block(x)
        return x
//  [python]
def predict(self, X_list, t_list, t_list_tr):
        
        with torch.no_grad():
            
            As = self._sample_At_star(t_list_tr, t_list, num_samples=0)[0]
            h_list = self._forward_latent_batch(X_list, t_list)
            pred_list = {}
            
            for tid, tt in enumerate(t_list):
                fid_t = self.dataset.func_t_to_fid(tt.item())
                ht = h_list[fid_t]
                At = As[:,:,tid]
                Nt = ht.shape[0]
                pred_t = torch.mm(ht, At.T).reshape([Nt, self.s_dim, self.s_dim])
                interp_pred_t_2d = \
                    t(
                        pred_t.unsqueeze(1), 
                        size=fid_t,
                        mode=self.interp,
                    ).squeeze(1)
                interp_pred_t_1d = interp_pred_t_2d.reshape([Nt, -1])
                pred_list[fid_t] = interp_pred_t_1d
            # 

            return pred_list
//  [python]
def predict(self, X_list, t_list, t_list_tr):
        
        with torch.no_grad():
            
            As = self._sample_At_star(t_list_tr, t_list, num_samples=0)[0]
            h_list = self._forward_latent_batch(X_list, t_list)
            pred_list = {}
            
            for tid, tt in enumerate(t_list):
                fid_t = self.dataset.func_t_to_fid(tt.item())
                ht = h_list[fid_t]
                At = As[:,:,tid]
                Nt = ht.shape[0]
                pred_t = torch.mm(ht, At.T).reshape([Nt, self.s_dim, self.s_dim])
                interp_pred_t_2d = \
                    torch.nn.functional.interpolate(
                        pred_t.unsqueeze(1), 
                        size=fid_t,
                        mode=self.interp,
                    ).squeeze(1)
                interp_pred_t_1d = interp_pred_t_2d.reshape([Nt, -1])
                pred_list[fid_t] = interp_pred_t_1d
            # 

            return pred_list
//  [python]
def forward(self, x, skip, concat=True):
        out = F.upsample_bilinear(
                x,
                size=(skip.size(2), skip.size(3)))
        if concat:
          out = torch.cat([out, skip], 1)
        return out
//  [python]
def forward(self, x, skip, concat=True):
        out = F.interpolate(
                x,
                size=(skip.size(2), skip.size(3)),
                mode="bilinear",
                align_corners=True)
        if concat:
          out = torch.cat([out, skip], 1)
        return out
//  [python]
def preprocess_single(video, resolution, sequence_length=None):
    # video: TCHW, {0, ..., 255} 
    video = video.float() / 255. # TCHW
    t, c, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = t(video, size=target_size, mode='bilinear',
                          align_corners=False)

    # center crop
    t, c, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
    video = video.permute(1, 0, 2, 3).contiguous() # CTHW

    video -= 0.5

    return video
//  [python]
def preprocess_single(video, resolution, sequence_length=None):
    # video: TCHW, {0, ..., 255} 
    video = video.float() / 255. # TCHW
    t, c, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = F.interpolate(video, size=target_size, mode='bilinear',
                          align_corners=False)

    # center crop
    t, c, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
    video = video.permute(1, 0, 2, 3).contiguous() # CTHW

    video -= 0.5

    return video
//  [python]
def upflow8(flow, mode='bilinear'):
    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
    return  8 * t.interpolate(flow, size=new_size, mode=mode, align_corners=False)
//  [python]
def upflow8(flow, mode='bilinear'):
    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
    return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=False)
//  [python]
def _eval_elbo_llh_t(self, ht, yt, tt):

        fid_t = self.dataset.func_t_to_fid(tt.item())
        tid = self.dataset.fid_list_tr.index(fid_t)
        
        N = yt.shape[0]
        
        yt_2d = yt.reshape([N, fid_t, fid_t]).unsqueeze(1)
        
        interp_yt_2d = \
            t.nn.functional.interpolate(
                yt_2d, 
                size=self.s_dim,
                mode=self.interp,
            ).squeeze(1)
        
        re_yt = interp_yt_2d.reshape([N, -1])
        
        term1 = torch.sum(torch.square(re_yt), dim=1, keepdim=True)
        EE_A_t = self.A.mu[:,:,tid]
        
        term2 = torch.bmm(ht.unsqueeze(1) @ (EE_A_t.T), re_yt.unsqueeze(-1)).squeeze(-1)
        
        L_d = torch.tril(self.A.chol_d)
        L_K = torch.tril(self.A.chol_K)
        L_T = torch.tril(self.A.chol_T)
        
        Sd = torch.mm(L_d, L_d.T)
        SK = torch.mm(L_K, L_K.T)
        ST = torch.mm(L_T, L_T.T)
        
        V = SK * ST[tid, tid]
        U = Sd
        
        EE_A_t_T_A_t = V * torch.trace(U) + torch.mm(EE_A_t.T, EE_A_t)
        
        term3 = torch.bmm((ht.unsqueeze(1))@EE_A_t_T_A_t, ht.unsqueeze(-1)).squeeze(-1)
        
        tau = torch.exp(self.log_tau)
        y_dim = self.s_dim**2
        
        llh_n = -0.5*y_dim*np.log(2*np.pi) + 0.5*y_dim*self.log_tau - \
            0.5*tau*(term1 - 2*term2 + term3)
        
        llh_t = torch.sum(llh_n)
        
        return llh_t
//  [python]
def _eval_elbo_llh_t(self, ht, yt, tt):

        fid_t = self.dataset.func_t_to_fid(tt.item())
        tid = self.dataset.fid_list_tr.index(fid_t)
        
        N = yt.shape[0]
        
        yt_2d = yt.reshape([N, fid_t, fid_t]).unsqueeze(1)
        
        interp_yt_2d = \
            torch.nn.functional.interpolate(
                yt_2d, 
                size=self.s_dim,
                mode=self.interp,
            ).squeeze(1)
        
        re_yt = interp_yt_2d.reshape([N, -1])
        
        term1 = torch.sum(torch.square(re_yt), dim=1, keepdim=True)
        EE_A_t = self.A.mu[:,:,tid]
        
        term2 = torch.bmm(ht.unsqueeze(1) @ (EE_A_t.T), re_yt.unsqueeze(-1)).squeeze(-1)
        
        L_d = torch.tril(self.A.chol_d)
        L_K = torch.tril(self.A.chol_K)
        L_T = torch.tril(self.A.chol_T)
        
        Sd = torch.mm(L_d, L_d.T)
        SK = torch.mm(L_K, L_K.T)
        ST = torch.mm(L_T, L_T.T)
        
        V = SK * ST[tid, tid]
        U = Sd
        
        EE_A_t_T_A_t = V * torch.trace(U) + torch.mm(EE_A_t.T, EE_A_t)
        
        term3 = torch.bmm((ht.unsqueeze(1))@EE_A_t_T_A_t, ht.unsqueeze(-1)).squeeze(-1)
        
        tau = torch.exp(self.log_tau)
        y_dim = self.s_dim**2
        
        llh_n = -0.5*y_dim*np.log(2*np.pi) + 0.5*y_dim*self.log_tau - \
            0.5*tau*(term1 - 2*term2 + term3)
        
        llh_t = torch.sum(llh_n)
        
        return llh_t
//  [python]
def depths_to_tensor(
        self, depths: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor in the range [-1, 1].
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        num_views = len(depths[0])
        for inner_list in depths:
            assert len(inner_list) == num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth) / self.max_depth
                tensor = tensor * 2 - 1
                tensor = t.upsample(
                    tensor[None, None],
                    scale_factor=self.image_size / tensor.shape[-1],
                    mode="nearest",
                )
                inner_batch.append(tensor.to(device=self.device, dtype=torch.float32))
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)
//  [python]
def depths_to_tensor(
        self, depths: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor in the range [-1, 1].
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        num_views = len(depths[0])
        for inner_list in depths:
            assert len(inner_list) == num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth) / self.max_depth
                tensor = tensor * 2 - 1
                tensor = F.interpolate(
                    tensor[None, None],
                    (self.image_size,) * 2,
                    mode="nearest",
                )
                inner_batch.append(tensor.to(device=self.device, dtype=torch.float32))
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)
//  [python]
def forward(self, x):
        if self.layer_type == 'none':
            return x
        else:
            return F.upsample(x, scale_factor=2, mode='nearest')
//  [python]
def forward(self, x):
        if self.layer_type == 'none':
            return x
        else:
            return F.interpolate(x, scale_factor=2, mode='nearest')
//  [python]
def thin_plate_dense(x1, y1, shape, step, lambd=.0, unroll_step_size=2**12):
    device = x1.device
    D, H, W = shape
    D1, H1, W1 = D//step, H//step, W//step
    
    x2 = F.affine_grid(torch.eye(3, 4, device=device).unsqueeze(0), (1, 1, D1, H1, W1), align_corners=True).view(-1, 3)
    tps = TPS()
    theta = tps.fit(x1[0], y1[0], lambd)
    
    y2 = torch.zeros((1, D1 * H1 * W1, 3), device=device)
    N = D1*H1*W1
    n = math.ceil(N/unroll_step_size)
    for j in range(n):
        j1 = j * unroll_step_size
        j2 = min((j + 1) * unroll_step_size, N)
        y2[0, j1:j2, :] = tps.z(x2[j1:j2], x1[0], theta)
        
    y2 = y2.view(1, D1, H1, W1, 3).permute(0, 4, 1, 2, 3)
    y2 = t(y2, (D, H, W), mode='trilinear', align_corners=True).permute(0, 2, 3, 4, 1)
    
    return y2
//  [python]
def thin_plate_dense(x1, y1, shape, step, lambd=.0, unroll_step_size=2**12):
    device = x1.device
    D, H, W = shape
    D1, H1, W1 = D//step, H//step, W//step
    
    x2 = F.affine_grid(torch.eye(3, 4, device=device).unsqueeze(0), (1, 1, D1, H1, W1), align_corners=True).view(-1, 3)
    tps = TPS()
    theta = tps.fit(x1[0], y1[0], lambd)
    
    y2 = torch.zeros((1, D1 * H1 * W1, 3), device=device)
    N = D1*H1*W1
    n = math.ceil(N/unroll_step_size)
    for j in range(n):
        j1 = j * unroll_step_size
        j2 = min((j + 1) * unroll_step_size, N)
        y2[0, j1:j2, :] = tps.z(x2[j1:j2], x1[0], theta)
        
    y2 = y2.view(1, D1, H1, W1, 3).permute(0, 4, 1, 2, 3)
    y2 = F.interpolate(y2, (D, H, W), mode='trilinear', align_corners=True).permute(0, 2, 3, 4, 1)
    
    return y2
//  [python]
def forward(self, x):
        xs = []
        x_sc = []
        
        for i in range(len(self.base)):
            x = self.base[i](x)
            if i in self.skip_nodes:
                xs.append(x)
        
        x = self.last_proj(x)
        x = self.last_pool(x)
        x2 = self.avg9x9(x)
        x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
        x = torch.cat([x,x2,x3],1)
        x = self.last_blk(x)
        
        for i in range(3):
            skip_x = xs[3-i]
            x = self.transUpBlocks[i](x, skip_x, (i<self.skip_lv))
            x = self.conv1x1_up[i](x)
            if self.SC[i] > 0:
              end = x.shape[1]
              x_sc.append( x[:,end-self.SC[i]:,:,:].contiguous() )
              x = x[:,:end-self.SC[i],:,:].contiguous()
            x2 = self.avg9x9(x)
            x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
            x = torch.cat([x,x2,x3],1)
            x = self.denseBlocksUp[i](x)

        scs = [x]
        for i in range(3):
          if self.SC[i] > 0:
            scs.insert(0,  t.interpolate(
                            x_sc[i], size=(x.size(2), x.size(3)), 
                            mode="bilinear", align_corners=True) )
        x = torch.cat(scs,1)
        x = self.conv_out(x)
        return x
//  [python]
def forward(self, x):
        xs = []
        x_sc = []
        
        for i in range(len(self.base)):
            x = self.base[i](x)
            if i in self.skip_nodes:
                xs.append(x)
        
        x = self.last_proj(x)
        x = self.last_pool(x)
        x2 = self.avg9x9(x)
        x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
        x = torch.cat([x,x2,x3],1)
        x = self.last_blk(x)
        
        for i in range(3):
            skip_x = xs[3-i]
            x = self.transUpBlocks[i](x, skip_x, (i<self.skip_lv))
            x = self.conv1x1_up[i](x)
            if self.SC[i] > 0:
              end = x.shape[1]
              x_sc.append( x[:,end-self.SC[i]:,:,:].contiguous() )
              x = x[:,:end-self.SC[i],:,:].contiguous()
            x2 = self.avg9x9(x)
            x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
            x = torch.cat([x,x2,x3],1)
            x = self.denseBlocksUp[i](x)

        scs = [x]
        for i in range(3):
          if self.SC[i] > 0:
            scs.insert(0,  F.interpolate(
                            x_sc[i], size=(x.size(2), x.size(3)), 
                            mode="bilinear", align_corners=True) )
        x = torch.cat(scs,1)
        x = self.conv_out(x)
        return x
//  [python]
def test_upsample_bilinear2d_output_size(self):
        def func(x: torch.Tensor):
            return torch.nn.functional.upsample_bilinear(x, size=(4, 4))

        onnx_program = torch.onnx.dynamo_export(func, torch.randn(1, 1, 2, 2))
        # If decomposition is skipped, the model will contain a Resize op instead of fine grained subgraph.
        assert_op_in_onnx_model(onnx_program.model_proto, "Resize")
//  [python]
def test_upsample_bilinear2d_output_size(self):
        def func(x: torch.Tensor):
            return torch.nn.functional.interpolate(x, size=(4, 4), mode="bilinear")

        onnx_program = torch.onnx.dynamo_export(func, torch.randn(1, 1, 2, 2))
        # If decomposition is skipped, the model will contain a Resize op instead of fine grained subgraph.
        assert_op_in_onnx_model(onnx_program.model_proto, "Resize")
//  [python]
def store_discr_outputs(self, batch):
        out_size = batch['image'].shape[2:]
        discr_real_out, _ = self.discriminator(batch['image'])
        discr_fake_out, _ = self.discriminator(batch['predicted_image'])
        batch['discr_output_real'] = t.interpolate(discr_real_out, size=out_size, mode='nearest')
        batch['discr_output_fake'] = t.interpolate(discr_fake_out, size=out_size, mode='nearest')
        batch['discr_output_diff'] = batch['discr_output_real'] - batch['discr_output_fake']
//  [python]
def store_discr_outputs(self, batch):
        out_size = batch['image'].shape[2:]
        discr_real_out, _ = self.discriminator(batch['image'])
        discr_fake_out, _ = self.discriminator(batch['predicted_image'])
        batch['discr_output_real'] = F.interpolate(discr_real_out, size=out_size, mode='nearest')
        batch['discr_output_fake'] = F.interpolate(discr_fake_out, size=out_size, mode='nearest')
        batch['discr_output_diff'] = batch['discr_output_real'] - batch['discr_output_fake']
//  [python]
def make_multiscale_noise(base_tensor, scales=6, scale_mode='bilinear'):
    batch_size, _, height, width = base_tensor.shape
    cur_height, cur_width = height, width
    result = []
    align_corners = False if scale_mode in ('bilinear', 'bicubic') else None
    for _ in range(scales):
        cur_sample = torch.randn(batch_size, 1, cur_height, cur_width, device=base_tensor.device)
        cur_sample_scaled = t.interpolate(cur_sample, size=(height, width), mode=scale_mode, align_corners=align_corners)
        result.append(cur_sample_scaled)
        cur_height //= 2
        cur_width //= 2
    return torch.cat(result, dim=1)
//  [python]
def make_multiscale_noise(base_tensor, scales=6, scale_mode='bilinear'):
    batch_size, _, height, width = base_tensor.shape
    cur_height, cur_width = height, width
    result = []
    align_corners = False if scale_mode in ('bilinear', 'bicubic') else None
    for _ in range(scales):
        cur_sample = torch.randn(batch_size, 1, cur_height, cur_width, device=base_tensor.device)
        cur_sample_scaled = F.interpolate(cur_sample, size=(height, width), mode=scale_mode, align_corners=align_corners)
        result.append(cur_sample_scaled)
        cur_height //= 2
        cur_width //= 2
    return torch.cat(result, dim=1)
//  [python]
def forward(self, x):
        """Forward function for upsampling operation.
        Args:
            x (torch.Tensor): Feature map to upsample.

        Returns:
            torch.Tensor: Feature map after upsampling.
        """
        x = F.upsample(x, scale_factor=2, mode='nearest')
        if self.with_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x):
        """Forward function for upsampling operation.
        Args:
            x (torch.Tensor): Feature map to upsample.

        Returns:
            torch.Tensor: Feature map after upsampling.
        """
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        if self.with_conv:
            x = self.conv(x)
        return x