jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > input_w:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def forward(self, x, segmap):
        normalized = self.param_free_norm(x)
        segmap = t.interpolate(segmap, size=x.size()[2:], mode='nearest')
        actv = self.mlp_shared(segmap)
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)
        out = normalized * (1 + gamma) + beta
        return out
//  [python]
def forward(self, x, segmap):
        normalized = self.param_free_norm(x)
        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
        actv = self.mlp_shared(segmap)
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)
        out = normalized * (1 + gamma) + beta
        return out
//  [python]
def test_interpolate(self):
        def _test_interpolate_helper(in_t, scale_factor, layer):
                out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                dim = len(in_t.shape) - 2
                out_shape = [1, 1] + [out_size] * dim
                out_t = m(in_t)
                self.assertEqual(torch.ones(out_shape), out_t)

                self.assertEqual(
                    t.interpolate(in_t, (out_size,) * dim, **kwargs),
                    t.interpolate(in_t, scale_factor=scale_factor, **kwargs))
                gradcheck(lambda x: t.interpolate(x, out_size, **kwargs), [in_t])
                gradgradcheck(lambda x: t.interpolate(x, out_size, **kwargs), [in_t])

        def _make_input(dim):
            size = [1, 1]
            size += [2] * dim
            return torch.ones(size, requires_grad=True)

        device_list = ['cpu']
        if TEST_CUDA:
            device_list.append('cuda')

        for device in device_list:
            for scale_factor in [0.5, 1.5, 2]:
                for mode in ['nearest', 'area']:
                    kwargs = dict(mode=mode)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    for input in [_make_input(1), _make_input(2), _make_input(3)]:
                        _test_interpolate_helper(input, scale_factor, m)

                for align_corners in [True, False]:
                    kwargs = dict(mode='linear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(1), scale_factor, m)

                    kwargs = dict(mode='bilinear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(2), scale_factor, m)

                    kwargs = dict(mode='trilinear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(3), scale_factor, m)
//  [python]
def test_interpolate(self):
        def _test_interpolate_helper(in_t, scale_factor, layer):
                out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                dim = len(in_t.shape) - 2
                out_shape = [1, 1] + [out_size] * dim
                out_t = m(in_t)
                self.assertEqual(torch.ones(out_shape), out_t)

                self.assertEqual(
                    F.interpolate(in_t, (out_size,) * dim, **kwargs),
                    F.interpolate(in_t, scale_factor=scale_factor, **kwargs))
                gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [in_t])
                gradgradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [in_t])

        def _make_input(dim):
            size = [1, 1]
            size += [2] * dim
            return torch.ones(size, requires_grad=True)

        device_list = ['cpu']
        if TEST_CUDA:
            device_list.append('cuda')

        for device in device_list:
            for scale_factor in [0.5, 1.5, 2]:
                for mode in ['nearest', 'area']:
                    kwargs = dict(mode=mode)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    for input in [_make_input(1), _make_input(2), _make_input(3)]:
                        _test_interpolate_helper(input, scale_factor, m)

                for align_corners in [True, False]:
                    kwargs = dict(mode='linear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(1), scale_factor, m)

                    kwargs = dict(mode='bilinear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(2), scale_factor, m)

                    kwargs = dict(mode='trilinear', align_corners=align_corners)
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs).to(device)
                    _test_interpolate_helper(_make_input(3), scale_factor, m)
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            inner_top_down = F.upsample(last_inner, scale_factor=2, mode="nearest")
            inner_lateral = getattr(self, inner_block)(feature)
            # TODO use size instead of scale to make it robust to different sizes
            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            # mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            results.insert(0, getattr(self, layer_block)(last_inner))

        if self.top_blocks is not None:
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        return tuple(results)
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")
            inner_lateral = getattr(self, inner_block)(feature)
            # TODO use size instead of scale to make it robust to different sizes
            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            # mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            results.insert(0, getattr(self, layer_block)(last_inner))

        if self.top_blocks is not None:
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        return tuple(results)
//  [python]
def score(self, features: Tensor, feature_shapes: tuple) -> Tensor:
        """Compute scores.

        Scores are either PCA-based feature reconstruction error (FRE) scores or
        the Gaussian density-based NLL scores

        Args:
            features (torch.Tensor): semantic features on which PCA and density modeling is performed.
            feature_shapes  (tuple): shape of `features` tensor. Used to generate anomaly map of correct shape.

        Returns:
            score (Tensor): numpy array of scores
        """
        feats_projected = self.pca_model.transform(features)
        if self.score_type == "nll":
            score = self.gaussian_model.score_samples(feats_projected)
        elif self.score_type == "fre":
            feats_reconstructed = self.pca_model.inverse_transform(feats_projected)
            fre = torch.square(features - feats_reconstructed).reshape(feature_shapes)
            fre_map = torch.unsqueeze(torch.sum(fre, dim=1), 1)
            score_map = t.interpolate(fre_map, size=self.input_size, mode="bilinear", align_corners=False)
            score = torch.sum(torch.square(features - feats_reconstructed), dim=1)
        else:
            raise ValueError(f"unsupported score type: {self.score_type}")

        if self.score_type == "nll":
            output = score
        else:
            output = score_map, score

        return output
//  [python]
def score(self, features: Tensor, feature_shapes: tuple) -> Tensor:
        """Compute scores.

        Scores are either PCA-based feature reconstruction error (FRE) scores or
        the Gaussian density-based NLL scores

        Args:
            features (torch.Tensor): semantic features on which PCA and density modeling is performed.
            feature_shapes  (tuple): shape of `features` tensor. Used to generate anomaly map of correct shape.

        Returns:
            score (Tensor): numpy array of scores
        """
        feats_projected = self.pca_model.transform(features)
        if self.score_type == "nll":
            score = self.gaussian_model.score_samples(feats_projected)
        elif self.score_type == "fre":
            feats_reconstructed = self.pca_model.inverse_transform(feats_projected)
            fre = torch.square(features - feats_reconstructed).reshape(feature_shapes)
            fre_map = torch.unsqueeze(torch.sum(fre, dim=1), 1)
            score_map = F.interpolate(fre_map, size=self.input_size, mode="bilinear", align_corners=False)
            score = torch.sum(torch.square(features - feats_reconstructed), dim=1)
        else:
            raise ValueError(f"unsupported score type: {self.score_type}")

        if self.score_type == "nll":
            output = score
        else:
            output = score_map, score

        return output
//  [python]
def forward(self, block_input: torch.Tensor):
        image_data = block_input.permute(0, 3, 1, 2)
        original_shape = image_data.shape[2:]
        image_data = torch.nn.functional.interpolate(image_data, (320, 320), mode='bilinear')
        image_data = (image_data / 255 - 0.485) / 0.229
        out = self.net(image_data)[0][:, 0:1]
        ma = torch.max(out)
        mi = torch.min(out)
        out = (out - mi) / (ma - mi) * 255
        out = torch.nn.functional.interpolate(out, original_shape, mode='bilinear')
        out = out[:, 0]
        out = out.to(dtype=torch.uint8, device=torch.device('cpu'), non_blocking=True).detach()
        return out
//  [python]
def imenlarge2(self, im):
        B, _, M, N = im.shape
        t1 = t(im, size=(int(4*M-3), int(4*N-3)), mode='bilinear', align_corners=False)
        t2 = torch.zeros([B, 1, 4*M-1, 4*N-1]).type(self.samplet.type())
        t2[:, :, 1: -1, 1:-1] = t1
        t2[:, :, 0, :] = 2*t2[:, :, 1, :] - t2[:, :, 2, :]
        t2[:, :, -1, :] = 2*t2[:, :, -2, :] - t2[:, :, -3, :]
        t2[:, :, :, 0] = 2*t2[:, :, :, 1] - t2[:, :, :, 2]
        t2[:, :, :, -1] = 2*t2[:, :, :, -2] - t2[:, :, :, -3]
        imu = t2[:, :, ::2, ::2]
        return imu
//  [python]
def imenlarge2(self, im):
        B, _, M, N = im.shape
        t1 = F.interpolate(im, size=(int(4*M-3), int(4*N-3)), mode='bilinear', align_corners=False)
        t2 = torch.zeros([B, 1, 4*M-1, 4*N-1]).type(self.samplet.type())
        t2[:, :, 1: -1, 1:-1] = t1
        t2[:, :, 0, :] = 2*t2[:, :, 1, :] - t2[:, :, 2, :]
        t2[:, :, -1, :] = 2*t2[:, :, -2, :] - t2[:, :, -3, :]
        t2[:, :, :, 0] = 2*t2[:, :, :, 1] - t2[:, :, :, 2]
        t2[:, :, :, -1] = 2*t2[:, :, :, -2] - t2[:, :, :, -3]
        imu = t2[:, :, ::2, ::2]
        return imu
//  [python]
def forward(self, x):
        #residual = x
        x3, x4, x5, x6 = self.backbone(x)

        outs = {}
        x = self.up0(x6, x5)
        x = self.up1(x, x4)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        x = t.upsample(x, scale_factor=4, mode='bilinear', align_corners=True)
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def forward(self, x):
        #residual = x
        x3, x4, x5, x6 = self.backbone(x)

        outs = {}
        x = self.up0(x6, x5)
        x = self.up1(x, x4)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        x = F.interpolate(x, scale_factor=4, align_corners=True, mode='bilinear')
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def forward(self, x, P2=None):
        # residual = x
        x3, x4, x5, x6 = self.backbone(x)

        outs = {}
        #x6 = F.relu(x6 + self.fullencoders(x6))
        x6 = self.fullencoders(x6)
        x = self.up0(x6, x5, P2=P2, scale=32)
        x = self.up1(x, x4, P2=P2, scale=16)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        #x = F.upsample_bilinear(x, scale_factor=4)
        x = F.upsample_bilinear(x, scale_factor=4)
        #x = self.up3(x)
        #x = self.up4(x, residual)
        #x = torch.cat([x, residual], dim=1)
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def forward(self, x, P2=None):
        # residual = x
        x3, x4, x5, x6 = self.backbone(x)

        outs = {}
        #x6 = F.relu(x6 + self.fullencoders(x6))
        x6 = self.fullencoders(x6)
        x = self.up0(x6, x5, P2=P2, scale=32)
        x = self.up1(x, x4, P2=P2, scale=16)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        #x = F.upsample_bilinear(x, scale_factor=4)
        x = F.interpolate(x, scale_factor=4, align_corners=True, mode='bilinear')
        #x = self.up3(x)
        #x = self.up4(x, residual)
        #x = torch.cat([x, residual], dim=1)
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def forward(self, tenInput):
        tenInput = tenInput * 255.0
        tenInput = tenInput - torch.tensor(data=[104.00698793, 116.66876762, 122.67891434], dtype=tenInput.dtype, device=tenInput.device).view(1, 3, 1, 1)

        tenVggOne = self.netVggOne(tenInput)
        tenVggTwo = self.netVggTwo(tenVggOne)
        tenVggThr = self.netVggThr(tenVggTwo)
        tenVggFou = self.netVggFou(tenVggThr)
        tenVggFiv = self.netVggFiv(tenVggFou)

        tenScoreOne = self.netScoreOne(tenVggOne)
        tenScoreTwo = self.netScoreTwo(tenVggTwo)
        tenScoreThr = self.netScoreThr(tenVggThr)
        tenScoreFou = self.netScoreFou(tenVggFou)
        tenScoreFiv = self.netScoreFiv(tenVggFiv)

        tenScoreOne = torch.nn.functional.upsample(input=tenScoreOne, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear')
        tenScoreTwo = torch.nn.functional.upsample(input=tenScoreTwo, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear')
        tenScoreThr = torch.nn.functional.upsample(input=tenScoreThr, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear')
        tenScoreFou = torch.nn.functional.upsample(input=tenScoreFou, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear')
        tenScoreFiv = torch.nn.functional.upsample(input=tenScoreFiv, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear')

        return self.netCombine(torch.cat([ tenScoreOne, tenScoreTwo, tenScoreThr, tenScoreFou, tenScoreFiv ], 1))
//  [python]
def forward(self, tenInput):
        tenInput = tenInput * 255.0
        tenInput = tenInput - torch.tensor(data=[104.00698793, 116.66876762, 122.67891434], dtype=tenInput.dtype, device=tenInput.device).view(1, 3, 1, 1)

        tenVggOne = self.netVggOne(tenInput)
        tenVggTwo = self.netVggTwo(tenVggOne)
        tenVggThr = self.netVggThr(tenVggTwo)
        tenVggFou = self.netVggFou(tenVggThr)
        tenVggFiv = self.netVggFiv(tenVggFou)

        tenScoreOne = self.netScoreOne(tenVggOne)
        tenScoreTwo = self.netScoreTwo(tenVggTwo)
        tenScoreThr = self.netScoreThr(tenVggThr)
        tenScoreFou = self.netScoreFou(tenVggFou)
        tenScoreFiv = self.netScoreFiv(tenVggFiv)

        tenScoreOne = torch.nn.functional.interpolate(input=tenScoreOne, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
        tenScoreTwo = torch.nn.functional.interpolate(input=tenScoreTwo, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
        tenScoreThr = torch.nn.functional.interpolate(input=tenScoreThr, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
        tenScoreFou = torch.nn.functional.interpolate(input=tenScoreFou, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
        tenScoreFiv = torch.nn.functional.interpolate(input=tenScoreFiv, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)

        return self.netCombine(torch.cat([ tenScoreOne, tenScoreTwo, tenScoreThr, tenScoreFou, tenScoreFiv ], 1))
//  [python]
def forward(self, tenFirst, tenSecond):
		intWidth = tenFirst.shape[3]
		intHeight = tenFirst.shape[2]

		intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
		intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))

		tenPreprocessedFirst = t.interpolate(input=tenFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
		tenPreprocessedSecond = t.interpolate(input=tenSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
		
		tenFirst = self.netExtractor(tenPreprocessedFirst)
		tenSecond = self.netExtractor(tenPreprocessedSecond)


		objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
		objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
		objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
		objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
		objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)

		tenFlow = objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])
		tenFlow = 20.0 * t.interpolate(input=tenFlow, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
		tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
		tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)

		return tenFlow
//  [python]
def forward(self, tenFirst, tenSecond):
		intWidth = tenFirst.shape[3]
		intHeight = tenFirst.shape[2]

		intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
		intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))

		tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
		tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
		
		tenFirst = self.netExtractor(tenPreprocessedFirst)
		tenSecond = self.netExtractor(tenPreprocessedSecond)


		objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
		objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
		objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
		objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
		objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)

		tenFlow = objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])
		tenFlow = 20.0 * torch.nn.functional.interpolate(input=tenFlow, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
		tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
		tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)

		return tenFlow
//  [python]
def repeat_expand_2d(content: torch.Tensor, target_len: int) -> torch.Tensor:
    # content : [h, t]
    src_len = content.shape[-1]
    if target_len < src_len:
        return content[:, :target_len]
    else:
        return torch.nn.functional.interpolate(
            content.unsqueeze(0), size=target_len, mode="nearest"
        ).squeeze(0)
//  [python]
def forward(self, x,y):
        input_size_h = x.size()[2]
        input_size_w = x.size()[3]

        x2 = t.interpolate(x, size=(int(input_size_h * 0.5), int(input_size_w * 0.5)), mode='bilinear',
                           align_corners=False)
        x3 = t.interpolate(x, size=(int(input_size_h * 1.5), int(input_size_w * 1.5)), mode='bilinear',
                           align_corners=False)
        x4 = t.interpolate(x, size=(int(input_size_h * 2), int(input_size_w * 2)), mode='bilinear',
                           align_corners=False)

        y2 = t.interpolate(y, size=(int(input_size_h * 0.5), int(input_size_w * 0.5)), mode='bilinear',
                           align_corners=False)
        y3 = t.interpolate(y, size=(int(input_size_h * 1.5), int(input_size_w * 1.5)), mode='bilinear',
                           align_corners=False)
        y4 = t.interpolate(y, size=(int(input_size_h * 2), int(input_size_w * 2)), mode='bilinear',
                           align_corners=False)

        x_cam, y_cam, x_mask, y_mask = super().forward(x,y)
        with torch.no_grad():
            x2, y2, x_mask2, y_mask2 = super().forward(x2, y2)
            x3, y3, x_mask3, y_mask3 = super().forward(x3, y3)
            x4, y4, x_mask4, y_mask4 = super().forward(x4, y4)

        x_mask2 = t.interpolate(x_mask2, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)
        x_mask3 = t.interpolate(x_mask3, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)
        x_mask4 = t.interpolate(x_mask4, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)

        y_mask2 = t.interpolate(y_mask2, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)
        y_mask3 = t.interpolate(y_mask3, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)
        y_mask4 = t.interpolate(y_mask4, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)

        cam_x = (F.relu(x_mask) + F.relu(x_mask2) +F.relu(x_mask3) +F.relu(x_mask4) )/4
        cam_y = (F.relu(y_mask) + F.relu(y_mask2) + F.relu(y_mask3) + F.relu(y_mask4)) / 4

        x = torchutils.gap2d(x_cam, keepdims=True)
        x = x.view(-1, 20)

        y = torchutils.gap2d(y_cam, keepdims=True)
        y = y.view(-1, 20)

        return x,y,x_cam,y_cam,cam_x,cam_y
//  [python]
def forward(self, x,y):
        input_size_h = x.size()[2]
        input_size_w = x.size()[3]

        x2 = F.interpolate(x, size=(int(input_size_h * 0.5), int(input_size_w * 0.5)), mode='bilinear',
                           align_corners=False)
        x3 = F.interpolate(x, size=(int(input_size_h * 1.5), int(input_size_w * 1.5)), mode='bilinear',
                           align_corners=False)
        x4 = F.interpolate(x, size=(int(input_size_h * 2), int(input_size_w * 2)), mode='bilinear',
                           align_corners=False)

        y2 = F.interpolate(y, size=(int(input_size_h * 0.5), int(input_size_w * 0.5)), mode='bilinear',
                           align_corners=False)
        y3 = F.interpolate(y, size=(int(input_size_h * 1.5), int(input_size_w * 1.5)), mode='bilinear',
                           align_corners=False)
        y4 = F.interpolate(y, size=(int(input_size_h * 2), int(input_size_w * 2)), mode='bilinear',
                           align_corners=False)

        x_cam, y_cam, x_mask, y_mask = super().forward(x,y)
        with torch.no_grad():
            x2, y2, x_mask2, y_mask2 = super().forward(x2, y2)
            x3, y3, x_mask3, y_mask3 = super().forward(x3, y3)
            x4, y4, x_mask4, y_mask4 = super().forward(x4, y4)

        x_mask2 = F.interpolate(x_mask2, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)
        x_mask3 = F.interpolate(x_mask3, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)
        x_mask4 = F.interpolate(x_mask4, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear', align_corners=False)

        y_mask2 = F.interpolate(y_mask2, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)
        y_mask3 = F.interpolate(y_mask3, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)
        y_mask4 = F.interpolate(y_mask4, size=(int(x_mask.shape[2]), int(x_mask.shape[3])), mode='bilinear',
                                align_corners=False)

        cam_x = (F.relu(x_mask) + F.relu(x_mask2) +F.relu(x_mask3) +F.relu(x_mask4) )/4
        cam_y = (F.relu(y_mask) + F.relu(y_mask2) + F.relu(y_mask3) + F.relu(y_mask4)) / 4

        x = torchutils.gap2d(x_cam, keepdims=True)
        x = x.view(-1, 20)

        y = torchutils.gap2d(y_cam, keepdims=True)
        y = y.view(-1, 20)

        return x,y,x_cam,y_cam,cam_x,cam_y
//  [python]
def clip_preprocess(image, size=224):
    mean = torch.tensor([ 0.48145466,0.4578275,0.40821073], device=image.device, dtype=image.dtype)
    std = torch.tensor([0.26862954,0.26130258,0.27577711], device=image.device, dtype=image.dtype)
    image = image.movedim(-1, 1)
    if not (image.shape[2] == size and image.shape[3] == size):
        scale = (size / min(image.shape[2], image.shape[3]))
        image = torch.nn.functional.interpolate(image, size=(round(scale * image.shape[2]), round(scale * image.shape[3])), mode="bicubic", antialias=True)
        h = (image.shape[2] - size)//2
        w = (image.shape[3] - size)//2
        image = image[:,:,h:h+size,w:w+size]
    image = torch.clip((255. * image), 0, 255).round() / 255.0
    return (image - mean.view([3,1,1])) / std.view([3,1,1])
//  [python]
def estimate(tenFirst, tenSecond):
	global netNetwork

	if netNetwork is None:
		netNetwork = Network().cuda().eval()
	# end

	assert(tenFirst.shape[1] == tenSecond.shape[1])
	assert(tenFirst.shape[2] == tenSecond.shape[2])

	intWidth = tenFirst.shape[2]
	intHeight = tenFirst.shape[1]

	assert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
	assert(intHeight == 436) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue

	tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth)
	tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth)

	intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
	intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))

	tenPreprocessedFirst = t.nn.functional.interpolate(input=tenPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
	tenPreprocessedSecond = t.nn.functional.interpolate(input=tenPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)

	tenFlow = 20.0 * t.nn.functional.interpolate(input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)

	tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
	tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)

	return tenFlow[0, :, :, :].cpu()
//  [python]
def estimate(tenFirst, tenSecond):
	global netNetwork

	if netNetwork is None:
		netNetwork = Network().cuda().eval()
	# end

	assert(tenFirst.shape[1] == tenSecond.shape[1])
	assert(tenFirst.shape[2] == tenSecond.shape[2])

	intWidth = tenFirst.shape[2]
	intHeight = tenFirst.shape[1]

	assert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
	assert(intHeight == 436) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue

	tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth)
	tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth)

	intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
	intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))

	tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
	tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)

	tenFlow = 20.0 * torch.nn.functional.interpolate(input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)

	tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
	tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)

	return tenFlow[0, :, :, :].cpu()
//  [python]
def resize_tensor(self, tensor):
        tensor = F.upsample_nearest(
            tensor.unsqueeze(0).float(), size=(1024, 2048)
        )
        return tensor[0]
//  [python]
def resize_tensor(self, tensor):
        tensor = F.interpolate(
            tensor.unsqueeze(0).float(), size=(1024, 2048), mode="nearest",
        )
        # tensor = F.upsample_nearest(
        #     tensor.unsqueeze(0).float(), size=(1024, 2048)
        # )
        return tensor[0]
//  [python]
def forward(self, pixel_values, position_embedding=None):
        embeddings = self.proj(pixel_values)

        if position_embedding is not None:
            # interpolate the position embedding to the corresponding size
            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1)
            position_embedding = position_embedding.permute(0, 3, 1, 2)
            patch_height, patch_width = embeddings.shape[2], embeddings.shape[3]
            position_embedding = F.upsample(position_embedding, size=(patch_height, patch_width), mode="bicubic")
            embeddings = embeddings + position_embedding

        embeddings = embeddings.flatten(2).transpose(1, 2)
        return embeddings
//  [python]
def forward(self, pixel_values, position_embedding=None):
        embeddings = self.proj(pixel_values)

        if position_embedding is not None:
            # interpolate the position embedding to the corresponding size
            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1)
            position_embedding = position_embedding.permute(0, 3, 1, 2)
            patch_height, patch_width = embeddings.shape[2], embeddings.shape[3]
            position_embedding = F.interpolate(position_embedding, size=(patch_height, patch_width), mode="bicubic")
            embeddings = embeddings + position_embedding

        embeddings = embeddings.flatten(2).transpose(1, 2)
        return embeddings
//  [python]
def test_upsample_bilinear(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return t.interpolate(x, scale_factor=2.0, mode='bilinear')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_bilinear(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, scale_factor=2.0, mode='bilinear')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_bilinear_align_corners(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return t.upsample_bilinear(x, scale_factor=2.0)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_bilinear_align_corners(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_nearest(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.upsample_nearest(x, scale_factor=2.0)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_nearest(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, scale_factor=2.0, mode='nearest')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_prelu_multi_channel_with_transposes(self):
        class Model(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.prelu = nn.PReLU(16)

            def forward(self, x):
                x = F.interpolate(x, scale_factor=2, mode='nearest')
                x = self.prelu(x)
                x = x.permute([0, 2, 3, 1])
                return x

        model = Model()
        model.eval()

        dummy_input = torch.randn(1, 16, 128, 1)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, check_stride=False)
//  [python]
def test_prelu_multi_channel_with_transposes(self):
        class Model(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.prelu = nn.PReLU(16)

            def forward(self, x):
                x = F.interpolate(x, scale_factor=2)
                x = self.prelu(x)
                x = x.permute([0, 2, 3, 1])
                return x

        model = Model()
        model.eval()

        dummy_input = torch.randn(1, 16, 128, 1)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, check_stride=False)
//  [python]
def test_upsample_bilinear_output_size(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.upsample_bilinear(x, size=(448, 448))

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_bilinear_output_size(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, size=(448, 448), mode='bilinear')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def AdamRegMIND(mind_fix,mind_mov,dense_flow):

    
    if(dense_flow.shape[-1]==3):
        dense_flow = dense_flow.permute(0,4,1,2,3)
    
    H,W,D = dense_flow[0,0].shape
    
    disp_hr = dense_flow.cuda().flip(1)*torch.tensor([H-1,W-1,D-1]).cuda().view(1,3,1,1,1)/2
    with torch.enable_grad(): 
        grid_sp = 2

       
        disp_lr = t.upsample(disp_hr,size=(H//grid_sp,W//grid_sp,D//grid_sp),mode='trilinear')
        net = nn.Sequential(nn.Conv3d(3,1,(H//grid_sp,W//grid_sp,D//grid_sp),bias=False))
        net[0].weight.data[:] = disp_lr.float().cpu().data/grid_sp
        net.cuda()
        optimizer = torch.optim.Adam(net.parameters(), lr=1)
        grid0 = F.affine_grid(torch.eye(3,4).unsqueeze(0).cuda(),(1,1,H//grid_sp,W//grid_sp,D//grid_sp),align_corners=False)
        #run Adam optimisation with diffusion regularisation and B-spline smoothing
        lambda_weight = .65
        for iter in range(50):
            optimizer.zero_grad()
            disp_sample = F.avg_pool3d(F.avg_pool3d(F.avg_pool3d(net[0].weight,3,stride=1,padding=1),3,stride=1,padding=1),\
                                       3,stride=1,padding=1).permute(0,2,3,4,1)
            reg_loss = lambda_weight*((disp_sample[0,:,1:,:]-disp_sample[0,:,:-1,:])**2).mean()+\
            lambda_weight*((disp_sample[0,1:,:,:]-disp_sample[0,:-1,:,:])**2).mean()+\
            lambda_weight*((disp_sample[0,:,:,1:]-disp_sample[0,:,:,:-1])**2).mean()
            scale = torch.tensor([(H//grid_sp-1)/2,(W//grid_sp-1)/2,(D//grid_sp-1)/2]).cuda().unsqueeze(0)
            grid_disp = grid0.view(-1,3).cuda().float()+((disp_sample.view(-1,3))/scale).flip(1).float()
            patch_mov_sampled = F.grid_sample(mind_mov.cuda().float(),grid_disp.view(1,H//grid_sp,W//grid_sp,D//grid_sp,3).cuda()\
                                              ,align_corners=False,mode='bilinear')
            sampled_cost = (patch_mov_sampled-mind_fix.cuda()).pow(2).mean(1)*12
            loss = sampled_cost.mean()
            (loss+reg_loss).backward()
            optimizer.step()
        fitted_grid = disp_sample.permute(0,4,1,2,3).detach()
        disp_hr = t.upsample(fitted_grid*grid_sp,size=(H,W,D),mode='trilinear')

    disp_smooth = F.avg_pool3d(F.avg_pool3d(F.avg_pool3d(disp_hr,3,padding=1,stride=1),3,padding=1,stride=1),3,padding=1,stride=1)


    disp_hr = torch.flip(disp_smooth/torch.tensor([H-1,W-1,D-1]).view(1,3,1,1,1).cuda()*2,[1])
    return disp_hr
//  [python]
def AdamRegMIND(mind_fix,mind_mov,dense_flow):

    
    if(dense_flow.shape[-1]==3):
        dense_flow = dense_flow.permute(0,4,1,2,3)
    
    H,W,D = dense_flow[0,0].shape
    
    disp_hr = dense_flow.cuda().flip(1)*torch.tensor([H-1,W-1,D-1]).cuda().view(1,3,1,1,1)/2
    with torch.enable_grad(): 
        grid_sp = 2

       
        disp_lr = F.interpolate(disp_hr,size=(H//grid_sp,W//grid_sp,D//grid_sp),mode='trilinear',align_corners=False)
        net = nn.Sequential(nn.Conv3d(3,1,(H//grid_sp,W//grid_sp,D//grid_sp),bias=False))
        net[0].weight.data[:] = disp_lr.float().cpu().data/grid_sp
        net.cuda()
        optimizer = torch.optim.Adam(net.parameters(), lr=1)
        grid0 = F.affine_grid(torch.eye(3,4).unsqueeze(0).cuda(),(1,1,H//grid_sp,W//grid_sp,D//grid_sp),align_corners=False)
        #run Adam optimisation with diffusion regularisation and B-spline smoothing
        lambda_weight = .65
        for iter in range(50):
            optimizer.zero_grad()
            disp_sample = F.avg_pool3d(F.avg_pool3d(F.avg_pool3d(net[0].weight,3,stride=1,padding=1),3,stride=1,padding=1),\
                                       3,stride=1,padding=1).permute(0,2,3,4,1)
            reg_loss = lambda_weight*((disp_sample[0,:,1:,:]-disp_sample[0,:,:-1,:])**2).mean()+\
            lambda_weight*((disp_sample[0,1:,:,:]-disp_sample[0,:-1,:,:])**2).mean()+\
            lambda_weight*((disp_sample[0,:,:,1:]-disp_sample[0,:,:,:-1])**2).mean()
            scale = torch.tensor([(H//grid_sp-1)/2,(W//grid_sp-1)/2,(D//grid_sp-1)/2]).cuda().unsqueeze(0)
            grid_disp = grid0.view(-1,3).cuda().float()+((disp_sample.view(-1,3))/scale).flip(1).float()
            patch_mov_sampled = F.grid_sample(mind_mov.cuda().float(),grid_disp.view(1,H//grid_sp,W//grid_sp,D//grid_sp,3).cuda()\
                                              ,align_corners=False,mode='bilinear')
            sampled_cost = (patch_mov_sampled-mind_fix.cuda()).pow(2).mean(1)*12
            loss = sampled_cost.mean()
            (loss+reg_loss).backward()
            optimizer.step()
        fitted_grid = disp_sample.permute(0,4,1,2,3).detach()
        disp_hr = F.interpolate(fitted_grid*grid_sp,size=(H,W,D),mode='trilinear',align_corners=False)

    disp_smooth = F.avg_pool3d(F.avg_pool3d(F.avg_pool3d(disp_hr,3,padding=1,stride=1),3,padding=1,stride=1),3,padding=1,stride=1)


    disp_hr = torch.flip(disp_smooth/torch.tensor([H-1,W-1,D-1]).view(1,3,1,1,1).cuda()*2,[1])
    return disp_hr
//  [python]
def _create_octave_mask(mask: Tensor, octave_shape: List[int]) -> Tensor:
    r"""Downsample a mask based on the given octave shape."""
    mask_shape = octave_shape[-2:]
    mask_octave = t.nn.functional.interpolate(mask, mask_shape, mode="bilinear", align_corners=False)
    return mask_octave.unsqueeze(1)
//  [python]
def _create_octave_mask(mask: Tensor, octave_shape: List[int]) -> Tensor:
    r"""Downsample a mask based on the given octave shape."""
    mask_shape = octave_shape[-2:]
    mask_octave = F.interpolate(mask, mask_shape, mode="bilinear", align_corners=False)
    return mask_octave.unsqueeze(1)
//  [python]
def forward(self, x):
        x = rearrange(x, "b c f h w -> b f c h w")
        x, ps = pack([x], "* c h w")
        x = t.interpolate(x, scale_factor=2.0, mode="nearest")
        x = self.conv(x)
        x = unpack(x, ps, "* c h w")[0]
        x = rearrange(x, "b f c h w -> b c f h w")
        return x
//  [python]
def forward(self, x):
        x = rearrange(x, "b c f h w -> b f c h w")
        x, ps = pack([x], "* c h w")
        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
        x = self.conv(x)
        x = unpack(x, ps, "* c h w")[0]
        x = rearrange(x, "b f c h w -> b c f h w")
        return x
//  [python]
def test_upsample_nearest_output_size(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return nn.Upsample(scale_factor=2, mode='nearest')(x)

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_nearest_output_size(self):
        dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, size=(448, 448), mode='nearest')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_nearest_output_size_complex(self):
        dummy_input = torch.randn(1, 3, 8, 22, dtype=torch.float32)

        def model(x):
            return torch.nn.functional.upsample_nearest(x, size=(15, 43))

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def test_upsample_nearest_output_size_complex(self):
        dummy_input = torch.randn(1, 3, 8, 22, dtype=torch.float32)

        def model(x):
            return F.interpolate(x, size=(15, 43), mode='nearest')

        model_path = get_model_path()
        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output)
//  [python]
def forward(self, x):
        x = rearrange(x, "b c f h w -> b c h w f")
        x, ps = pack([x], "b * f")
        if x.size(-1) > 1:
            x = torch.concat((x[:,:,:1], t.interpolate(x[:,:,1:], scale_factor=2.0, mode="linear")), dim=-1)
        else:
            x = x
        x = unpack(x, ps, "b * f")[0]
        x = rearrange(x, "b c h w f -> b c f h w")
        x = self.conv(x)
        return x
//  [python]
def forward(self, x):
        x = rearrange(x, "b c f h w -> b c h w f")
        x, ps = pack([x], "b * f")
        if x.size(-1) > 1:
            x = torch.concat((x[:,:,:1], F.interpolate(x[:,:,1:], scale_factor=2.0, mode="linear")), dim=-1)
        else:
            x = x
        x = unpack(x, ps, "b * f")[0]
        x = rearrange(x, "b c h w f -> b c f h w")
        x = self.conv(x)
        return x
//  [python]
def forward(self, x, reference=None):
        if self.mode == 'upsample':
            assert reference is not None
            x = F.interpolate(x, size=reference.shape[-2:], mode='nearest')
        elif self.mode == 'downsample':
            x = F.adaptive_avg_pool2d(
                x, output_size=tuple(u // 2 for u in x.shape[-2:]))
        return x
//  [python]
def forward(self, x: Tensor, out_shape: list[int]) -> Tensor:
        out = self.features(x)
        if out.shape[-2:] != out_shape:
            out = F.upsample(out, size=out_shape, mode="bilinear")
        return out
//  [python]
def forward(self, x: Tensor, out_shape: list[int]) -> Tensor:
        out = self.features(x)
        if out.shape[-2:] != out_shape:
            out = F.interpolate(out, out_shape, mode="bilinear")
        return out