jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x, gts=None):

        x2, x3, x4, x5 = self.backbone(x)

        outs = {}

        metric = self.mlayer(x5)

        x = self.up_4(x5, x4)
        
        d = self.vlayer(x)

        if self.training:
            var_loss = self.var_loss(x, d, gts)


        x, d  = self.ref_4(x, d)

        d_u4 = F.interpolate(d, scale_factor=16, mode='bilinear', align_corners=True)

        x = self.up_3(x, x3)

        x, d = self.ref_3(x, F.interpolate(d, scale_factor=2, mode='bilinear', align_corners=True))

        d_u3 = F.interpolate(d, scale_factor=8, mode='bilinear', align_corners=True)

        x = self.up_2(x, x2)

        x, d = self.ref_2(x, F.interpolate(d, scale_factor=2, mode='bilinear', align_corners=True))

        d_u2 = F.interpolate(d, scale_factor=4, mode='bilinear', align_corners=True)

        d = d_u2 + d_u3 + d_u4

        d = torch.sigmoid(metric[:, 0:1]) * (self.outc(d) + torch.exp(metric[:, 1:2]))

        outs['scale_1'] = d

        if self.training:
            si_loss = self.si_loss(outs, gts)
            return outs['scale_1'], var_loss + si_loss
        else:
            return outs['scale_1']
//  [python]
def test_interpolate(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest2d and upsample_bilinear2d
        """
        X, (scale, zero_point, torch_type) = X

        if scale_factor is not None:
            size = None
        if mode in ("nearest", "nearest-exact"):
            align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 3, 1, 2])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 3, 1, 2])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)

        X_ref = t.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": t.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }
        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={} X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def test_interpolate(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest2d and upsample_bilinear2d
        """
        X, (scale, zero_point, torch_type) = X

        if scale_factor is not None:
            size = None
        if mode in ("nearest", "nearest-exact"):
            align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 3, 1, 2])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 3, 1, 2])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)

        X_ref = torch.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": torch.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }
        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={} X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def upsample_bilinear(self,x,scale):
        scale_int = math.ceil(scale)
        N,C,H,W = x.size()
        x = t.interpolate(x, size=(round(H*scale), round(W*scale)), mode='bicubic')
        _,_,H2,W2 = x.size()
        # x = torch.stack([x]*2,2)
        # x = torch.stack([x]*2,4)
        # x = x.view(N,C,H2,1,W2,1)
        x = torch.cat([x]*2,2)
        x = torch.cat([x]*2,3)

        return x.contiguous().view(N,C,H2*2,W2*2)[:,:,:(H*scale_int), :(W*scale_int)]
//  [python]
def upsample_bilinear(self,x,scale):
        scale_int = math.ceil(scale)
        N,C,H,W = x.size()
        x = F.interpolate(x, size=(round(H*scale), round(W*scale)), mode='bicubic')
        _,_,H2,W2 = x.size()
        # x = torch.stack([x]*2,2)
        # x = torch.stack([x]*2,4)
        # x = x.view(N,C,H2,1,W2,1)
        x = torch.cat([x]*2,2)
        x = torch.cat([x]*2,3)

        return x.contiguous().view(N,C,H2*2,W2*2)[:,:,:(H*scale_int), :(W*scale_int)]
//  [python]
def test_interpolate3d(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest3d
        """
        X, (scale, zero_point, torch_type) = X
        if scale_factor is not None:
            size = None

        align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 4, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 4, 1, 2, 3])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 4, 1, 2, 3])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)
        X_ref = t.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": t.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }

        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={}, X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def test_interpolate3d(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest3d
        """
        X, (scale, zero_point, torch_type) = X
        if scale_factor is not None:
            size = None

        align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 4, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 4, 1, 2, 3])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 4, 1, 2, 3])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)
        X_ref = torch.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": torch.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }

        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={}, X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def apply_conv(self, source, guidance, conv, activation):
        big_source = t.interpolate(source, scale_factor=2, mode="bilinear")
        _, _, h, w = big_source.shape
        small_guidance = t.interpolate(guidance, (h, w), mode="bilinear")
        output = activation(conv(torch.cat([big_source, small_guidance], dim=1)))
        return big_source + output
//  [python]
def apply_conv(self, source, guidance, conv, activation):
        big_source = F.interpolate(source, scale_factor=2, mode="bilinear")
        _, _, h, w = big_source.shape
        small_guidance = F.interpolate(guidance, (h, w), mode="bilinear")
        output = activation(conv(torch.cat([big_source, small_guidance], dim=1)))
        return big_source + output
//  [python]
def forward(self, feats, img):
        _, _, h, w = img.shape
        return F.interpolate(feats, (h, w), mode="bilinear", align_corners=False)
//  [python]
def forward(self, feats, img):
        _, _, h, w = img.shape
        return F.interpolate(feats, (h, w), mode="bilinear")
//  [python]
def resize_image(img, size):
    return torch.nn.functional.interpolate(img, (size, size), mode='bilinear')
//  [python]
def forward(self, x, skip=None):
        x = F.upsample(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def forward(self, x, skip=None):
        x = F.interpolate(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def forward(self, x):

        e1 = self.Conv0(x)
        e2 = self.Conv1(e1)
        e2 = self.Conv2(e2)
        
        e3 = self.Maxpool3(e2)
        e3_1 = self.Conv3(e3)

        e4 = self.Maxpool4(e3_1)
        e4_1 = self.Conv4(e4)

        e5 = self.Maxpool5(e4_1)
        e5_1 = self.Conv5(e5)
        
        e6 = self.Maxpool6(e5_1)
        e6_1 = self.Conv6(e6)
        
        e7 = self.Conv7(e6_1)
        e8 = self.Conv8(e7)
        
        d1 =  t.upsample(e8, scale_factor=2, mode='nearest')
        cat0 = torch.cat([d1,e5_1],1)
        
        e9  = self.Conv9(cat0)
        e10 = self.Conv10(e9)
        
        d2 =  t.upsample(e10, scale_factor=2, mode='nearest')
        cat1 = torch.cat([d2,e4_1],1)
        
        e11 = self.Conv11(cat1)
        e12 = self.Conv12(e11)
        
        d3 =  t.upsample(e12, scale_factor=2, mode='nearest')
        cat2 = torch.cat([d3,e3_1],1)
        
        e13 = self.Conv13(cat2)
        e14 = self.Conv14(e13)
        
        d4 =  t.upsample(e14, scale_factor=2, mode='nearest')
        cat3 = torch.cat([d4,e2],1)
        
        e15 = self.Conv15(cat3)
        e16 = self.Conv16(e15)
        
        out = self.Conv_out(e16)
        out = self.active(out)
        return out
//  [python]
def forward(self, x):

        e1 = self.Conv0(x)
        e2 = self.Conv1(e1)
        e2 = self.Conv2(e2)
        
        e3 = self.Maxpool3(e2)
        e3_1 = self.Conv3(e3)

        e4 = self.Maxpool4(e3_1)
        e4_1 = self.Conv4(e4)

        e5 = self.Maxpool5(e4_1)
        e5_1 = self.Conv5(e5)
        
        e6 = self.Maxpool6(e5_1)
        e6_1 = self.Conv6(e6)
        
        e7 = self.Conv7(e6_1)
        e8 = self.Conv8(e7)
        
        d1 =  nn.functional.interpolate(e8, scale_factor=2)
        cat0 = torch.cat([d1,e5_1],1)
        
        e9  = self.Conv9(cat0)
        e10 = self.Conv10(e9)
        
        d2 =  nn.functional.interpolate(e10, scale_factor=2)
        cat1 = torch.cat([d2,e4_1],1)
        
        e11 = self.Conv11(cat1)
        e12 = self.Conv12(e11)
        
        d3 =  nn.functional.interpolate(e12, scale_factor=2)
        cat2 = torch.cat([d3,e3_1],1)
        
        e13 = self.Conv13(cat2)
        e14 = self.Conv14(e13)
        
        d4 =  nn.functional.interpolate(e14, scale_factor=2)
        cat3 = torch.cat([d4,e2],1)
        
        e15 = self.Conv15(cat3)
        e16 = self.Conv16(e15)
        
        out = self.Conv_out(e16)
        out = self.active(out)
        return out
//  [python]
def forward(self, source, guidance):
        b, c, h, w = source.shape
        up_source = t.interpolate(source, (h * 2, w * 2), mode="nearest")
        assert h == w
        lr_cord = torch.linspace(0, h, steps=h, device=source.device)
        hr_cord = torch.linspace(0, h, steps=2 * h, device=source.device)
        lr_coords = torch.cat([x.unsqueeze(0) for x in torch.meshgrid(lr_cord, lr_cord)], dim=0).unsqueeze(0)
        hr_coords = torch.cat([x.unsqueeze(0) for x in torch.meshgrid(hr_cord, hr_cord)], dim=0).unsqueeze(0)
        up_lr_coords = t.interpolate(lr_coords, (h * 2, w * 2), mode="nearest")
        coord_diff = up_lr_coords - hr_coords
        coord_diff_feats = self.sin_feats(coord_diff)
        c2 = coord_diff_feats.shape[1]
        bcast_coord_feats = torch.broadcast_to(coord_diff_feats, (b, c2, h * 2, w * 2))
        return self.mlp(torch.cat([up_source, bcast_coord_feats], dim=1))
//  [python]
def forward(self, source, guidance):
        b, c, h, w = source.shape
        up_source = F.interpolate(source, (h * 2, w * 2), mode="nearest")
        assert h == w
        lr_cord = torch.linspace(0, h, steps=h, device=source.device)
        hr_cord = torch.linspace(0, h, steps=2 * h, device=source.device)
        lr_coords = torch.cat([x.unsqueeze(0) for x in torch.meshgrid(lr_cord, lr_cord)], dim=0).unsqueeze(0)
        hr_coords = torch.cat([x.unsqueeze(0) for x in torch.meshgrid(hr_cord, hr_cord)], dim=0).unsqueeze(0)
        up_lr_coords = F.interpolate(lr_coords, (h * 2, w * 2), mode="nearest")
        coord_diff = up_lr_coords - hr_coords
        coord_diff_feats = self.sin_feats(coord_diff)
        c2 = coord_diff_feats.shape[1]
        bcast_coord_feats = torch.broadcast_to(coord_diff_feats, (b, c2, h * 2, w * 2))
        return self.mlp(torch.cat([up_source, bcast_coord_feats], dim=1))
//  [python]
def forward(self, x):
        size = x.size()[2:]
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        outputs = list()
        x = self.head(x1, x2, x3, x4, x5)
        x = F.upsample(x, size, mode='bilinear', align_corners=True)

        outputs.append(x)
        return tuple(outputs)
//  [python]
def forward(self, x):
        size = x.size()[2:]
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        outputs = list()
        x = self.head(x1, x2, x3, x4, x5)
        x = F.interpolate(x, size, mode='bilinear', align_corners=True)

        outputs.append(x)
        return tuple(outputs)
//  [python]
def scale_tensor(img_tensor, default_scale=256):
    _, _, w, h = img_tensor.shape
    if w < h:
        ow = default_scale
        oh = h / w * default_scale
    else:
        oh = default_scale
        ow = w / h * default_scale

    oh = int(round(oh / 16) * 16)
    ow = int(round(ow / 16) * 16)

    return t.interpolate(img_tensor, [ow, oh], mode="bilinear")
//  [python]
def scale_tensor(img_tensor, default_scale=256):
    _, _, w, h = img_tensor.shape
    if w < h:
        ow = default_scale
        oh = h / w * default_scale
    else:
        oh = default_scale
        ow = w / h * default_scale

    oh = int(round(oh / 16) * 16)
    ow = int(round(ow / 16) * 16)

    return F.interpolate(img_tensor, [ow, oh], mode="bilinear")
//  [python]
def deform_input(self, inp, deformation):
        _, d_old, h_old, w_old, _ = deformation.shape
        _, _, d, h, w = inp.shape
        if d_old != d or h_old != h or w_old != w:
            deformation = deformation.permute(0, 4, 1, 2, 3)
            deformation = t.interpolate(deformation, size=(d, h, w), mode="trilinear")
            deformation = deformation.permute(0, 2, 3, 4, 1)
        return F.grid_sample(inp, deformation)
//  [python]
def deform_input(self, inp, deformation):
        _, d_old, h_old, w_old, _ = deformation.shape
        _, _, d, h, w = inp.shape
        if d_old != d or h_old != h or w_old != w:
            deformation = deformation.permute(0, 4, 1, 2, 3)
            deformation = F.interpolate(deformation, size=(d, h, w), mode="trilinear")
            deformation = deformation.permute(0, 2, 3, 4, 1)
        return F.grid_sample(inp, deformation)
//  [python]
def forward(self, batched_inputs):
        """
        Args:
            batched_inputs (dict): expected keys: "img", Optional["caption"]

        """

        features = []

        image = batched_inputs["img"]
        batch_size = image.shape[0]

        if self.image_preprocess is None:
            normalized_image = (image - self.pixel_mean) / self.pixel_std
        else:
            normalized_image = self.image_preprocess((image - self.pixel_mean) / self.pixel_std)

        if "caption" in batched_inputs:
            captions = batched_inputs["caption"]
        else:
            captions = [""] * batch_size

        # latent_image = self.ldm.encode_to_latent(normalized_image)
        latent_image, encoder_features = self.encode_to_latent(normalized_image)
        cond_inputs = batched_inputs.get("cond_inputs", self.ldm.embed_text(captions))

        unet_features = []
        for i, t in enumerate(self.steps):

            if "cond_emb" in batched_inputs:
                cond_emb = batched_inputs["cond_emb"][:, i]
            else:
                cond_emb = None

            if t < 0:
                noisy_latent_image = latent_image
                # use 0 as no noise timestep
                t = torch.tensor([0], device=self.device).expand(batch_size)
            else:
                t = torch.tensor([t], device=self.device).expand(batch_size)
                if self.shared_noise is not None:
                    if self.shared_noise.shape[2:] != latent_image.shape[2:]:
                        assert self.image_preprocess is None
                        shared_noise = F.interpolate(
                            self.shared_noise,
                            size=latent_image.shape[2:],
                            mode="bicubic",
                            align_corners=False,
                        )
                    else:
                        shared_noise = self.shared_noise
                    noise = shared_noise.expand_as(latent_image)
                else:
                    noise = None

                noisy_latent_image = self.ldm.diffusion.q_sample(latent_image, t, noise)
            # self.ldm.ldm.apply_model(noisy_latent_image, t, cond_inputs)
            _, cond_unet_features = self.unet_forward(
                noisy_latent_image, t, cond_inputs, cond_emb=cond_emb
            )
            unet_features.extend(cond_unet_features)

        # self.ldm.decode_from_latent(latent_image)
        _, decoder_features = self.decode_to_image(latent_image)

        features = [*encoder_features, *unet_features, *decoder_features]

        assert len(features) == len(
            self.feature_dims
        ), f"{len(features)} != {len(self.feature_dims)}"

        for indices in self.grouped_indices:
            for idx in indices:
                if self.image_preprocess is not None:
                    continue
                assert image.shape[-2] // self.feature_strides[idx] == features[idx].shape[-2]
                assert image.shape[-1] // self.feature_strides[idx] == features[idx].shape[-1]

        return features
//  [python]
def forward(self, x, trimapmask, sampleidx=None):
        # set outputs
        outs = []

        # get outs[0]
        outs.append(self.shortcut[0](x))

        # forward patch-embedding layer
        x = self.patch_embed(x)
        _, _, Wh, Ww = x.shape
        trimapmask = t.interpolate(trimapmask, scale_factor=1/4, mode='nearest')

        # get outs[1]
        outs.append(self.shortcut[1](t.interpolate(x, scale_factor=2.0, mode='bilinear')))

        # dropout
        x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        # get outs[2~5]
        for i in range(self.num_layers):
            layer = self.layers[i]

            trimapmask_ = t.interpolate(trimapmask, scale_factor=1/(pow(2, i)), mode='nearest')
            area_fg = trimapmask_[:, 0, :, :].unsqueeze(-1)  # background area
            area_bg = trimapmask_[:, 2, :, :].unsqueeze(-1)  # foreground area
            area_uk = trimapmask_[:, 1, :, :].unsqueeze(-1)  # unknown area

            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww, area_fg=area_fg, area_bg=area_bg, area_uk=area_uk)

            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                x_out = norm_layer(x_out)
                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
                out = self.shortcut[i+2](out)
                outs.append(out)

        return tuple(outs)
//  [python]
def forward(self, x, trimapmask, sampleidx=None):
        # set outputs
        outs = []

        # get outs[0]
        outs.append(self.shortcut[0](x))

        # forward patch-embedding layer
        x = self.patch_embed(x)
        _, _, Wh, Ww = x.shape
        trimapmask = F.interpolate(trimapmask, scale_factor=1/4, mode='nearest')

        # get outs[1]
        outs.append(self.shortcut[1](F.upsample_bilinear(x, scale_factor=2.0)))

        # dropout
        x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        # get outs[2~5]
        for i in range(self.num_layers):
            layer = self.layers[i]

            trimapmask_ = F.interpolate(trimapmask, scale_factor=1/(pow(2, i)), mode='nearest')
            area_fg = trimapmask_[:, 0, :, :].unsqueeze(-1)  # background area
            area_bg = trimapmask_[:, 2, :, :].unsqueeze(-1)  # foreground area
            area_uk = trimapmask_[:, 1, :, :].unsqueeze(-1)  # unknown area

            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww, area_fg=area_fg, area_bg=area_bg, area_uk=area_uk)

            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                x_out = norm_layer(x_out)
                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
                out = self.shortcut[i+2](out)
                outs.append(out)

        return tuple(outs)
//  [python]
def scale_volume_grid(self, new_world_size):
        if self.channels == 0:
            self.grid = nn.Parameter(
                torch.zeros([1, self.channels, *new_world_size]))
        else:
            self.grid = nn.Parameter(
                t.interpolate(
                    self.grid.data,
                    size=tuple(new_world_size),
                    mode='trilinear',
                    align_corners=True))
//  [python]
def scale_volume_grid(self, new_world_size):
        if self.channels == 0:
            self.grid = nn.Parameter(
                torch.zeros([1, self.channels, *new_world_size]))
        else:
            self.grid = nn.Parameter(
                F.interpolate(
                    self.grid.data,
                    size=tuple(new_world_size),
                    mode='trilinear',
                    align_corners=True))
//  [python]
def forward(self, x):
        x, x_hw = x[0], x[1]
        x = self.linear_fuse(x)
        x = self.dropout(x)
        x = self.conv_seg(x)
        if self.upsample == 'intepolate' or self.training or \
                self.num_classes < 30:
            x = F.upsample(
                x, size=x_hw, mode='bilinear', align_corners=self.align_corners)
        elif self.upsample == 'vim':
            labelset = torch.unique(torch.argmax(x, 1))
            x = torch.gather(x, 1, labelset)
            x = F.upsample(
                x, size=x_hw, mode='bilinear', align_corners=self.align_corners)

            pred = torch.argmax(x, 1)
            pred_retrieve = torch.zeros(pred.shape, dtype=torch.int32)
            for i, val in enumerate(labelset):
                pred_retrieve[pred == i] = labelset[i].cast('int32')

            x = pred_retrieve
        else:
            raise NotImplementedError(self.upsample, ' is not implemented')

        return [x]
//  [python]
def forward(self, x):
        x, x_hw = x[0], x[1]
        x = self.linear_fuse(x)
        x = self.dropout(x)
        x = self.conv_seg(x)
        if self.upsample == 'intepolate' or self.training or \
                self.num_classes < 30:
            x = F.interpolate(
                x, x_hw, mode='bilinear', align_corners=self.align_corners)
        elif self.upsample == 'vim':
            labelset = torch.unique(torch.argmax(x, 1))
            x = torch.gather(x, 1, labelset)
            x = F.interpolate(
                x, x_hw, mode='bilinear', align_corners=self.align_corners)

            pred = torch.argmax(x, 1)
            pred_retrieve = torch.zeros(pred.shape, dtype=torch.int32)
            for i, val in enumerate(labelset):
                pred_retrieve[pred == i] = labelset[i].cast('int32')

            x = pred_retrieve
        else:
            raise NotImplementedError(self.upsample, ' is not implemented')

        return [x]
//  [python]
def apply_image(self, img, interp=None):
        assert img.shape[:2] == (self.h, self.w)
        assert len(img.shape) <= 4
        interp_method = interp if interp is not None else self.interp

        if img.dtype == np.uint8:
            if len(img.shape) > 2 and img.shape[2] == 1:
                pil_image = Image.fromarray(img[:, :, 0], mode='L')
            else:
                pil_image = Image.fromarray(img)
            pil_image = pil_image.resize((self.new_w, self.new_h),
                                         interp_method)
            ret = np.asarray(pil_image)
            if len(img.shape) > 2 and img.shape[2] == 1:
                ret = np.expand_dims(ret, -1)
        else:
            # PIL only supports uint8
            if any(x < 0 for x in img.strides):
                img = np.ascontiguousarray(img)
            img = torch.from_numpy(img)
            shape = list(img.shape)
            shape_4d = shape[:2] + [1] * (4 - len(shape)) + shape[2:]
            img = img.view(shape_4d).permute(2, 3, 0, 1)  # hw(c) -> nchw
            _PIL_RESIZE_TO_INTERPOLATE_MODE = {
                Image.NEAREST: 'nearest',
                Image.BILINEAR: 'bilinear',
                Image.BICUBIC: 'bicubic',
            }
            mode = _PIL_RESIZE_TO_INTERPOLATE_MODE[interp_method]
            align_corners = None if mode == 'nearest' else False
            img = t.interpolate(
                img, (self.new_h, self.new_w),
                mode=mode,
                align_corners=align_corners)
            shape[:2] = (self.new_h, self.new_w)
            ret = img.permute(2, 3, 0, 1).view(shape).numpy()  # nchw -> hw(c)

        return ret
//  [python]
def apply_image(self, img, interp=None):
        assert img.shape[:2] == (self.h, self.w)
        assert len(img.shape) <= 4
        interp_method = interp if interp is not None else self.interp

        if img.dtype == np.uint8:
            if len(img.shape) > 2 and img.shape[2] == 1:
                pil_image = Image.fromarray(img[:, :, 0], mode='L')
            else:
                pil_image = Image.fromarray(img)
            pil_image = pil_image.resize((self.new_w, self.new_h),
                                         interp_method)
            ret = np.asarray(pil_image)
            if len(img.shape) > 2 and img.shape[2] == 1:
                ret = np.expand_dims(ret, -1)
        else:
            # PIL only supports uint8
            if any(x < 0 for x in img.strides):
                img = np.ascontiguousarray(img)
            img = torch.from_numpy(img)
            shape = list(img.shape)
            shape_4d = shape[:2] + [1] * (4 - len(shape)) + shape[2:]
            img = img.view(shape_4d).permute(2, 3, 0, 1)  # hw(c) -> nchw
            _PIL_RESIZE_TO_INTERPOLATE_MODE = {
                Image.NEAREST: 'nearest',
                Image.BILINEAR: 'bilinear',
                Image.BICUBIC: 'bicubic',
            }
            mode = _PIL_RESIZE_TO_INTERPOLATE_MODE[interp_method]
            align_corners = None if mode == 'nearest' else False
            img = F.interpolate(
                img, (self.new_h, self.new_w),
                mode=mode,
                align_corners=align_corners)
            shape[:2] = (self.new_h, self.new_w)
            ret = img.permute(2, 3, 0, 1).view(shape).numpy()  # nchw -> hw(c)

        return ret
//  [python]
def forward(self, source_image, kp_driving, kp_source):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)
        out = self.second(out)
        bs, c, h, w = out.shape
        # print(out.shape)
        feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h, w)
        feature_3d = self.resblocks_3d(feature_3d)

        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source)
            output_dict["mask"] = dense_motion["mask"]

            if "occlusion_map" in dense_motion:
                occlusion_map = dense_motion["occlusion_map"]
                output_dict["occlusion_map"] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion["deformation"]
            out = self.deform_input(feature_3d, deformation)

            bs, c, d, h, w = out.shape
            out = out.view(bs, c * d, h, w)
            out = self.third(out)
            out = self.fourth(out)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = t.interpolate(occlusion_map, size=out.shape[2:], mode="bilinear")
                out = out * occlusion_map

        # Decoding part
        out = self.resblocks_2d(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def forward(self, source_image, kp_driving, kp_source):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)
        out = self.second(out)
        bs, c, h, w = out.shape
        # print(out.shape)
        feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h, w)
        feature_3d = self.resblocks_3d(feature_3d)

        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source)
            output_dict["mask"] = dense_motion["mask"]

            if "occlusion_map" in dense_motion:
                occlusion_map = dense_motion["occlusion_map"]
                output_dict["occlusion_map"] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion["deformation"]
            out = self.deform_input(feature_3d, deformation)

            bs, c, d, h, w = out.shape
            out = out.view(bs, c * d, h, w)
            out = self.third(out)
            out = self.fourth(out)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode="bilinear")
                out = out * occlusion_map

        # Decoding part
        out = self.resblocks_2d(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def get_mask(self, maps, alpha, use_pool):
        k = 1
        maps = (maps * alpha).sum(-1).mean(1)
        if use_pool:
            maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k))
        mask = t.nn.functional.interpolate(maps, size=(64, 64))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.th[1 - int(use_pool)])
        mask = mask[:1] + mask
        return mask
//  [python]
def get_mask(self, maps, alpha, use_pool):
        k = 1
        maps = (maps * alpha).sum(-1).mean(1)
        if use_pool:
            maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k))
        mask = nnf.interpolate(maps, size=(64, 64))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.th[1 - int(use_pool)])
        mask = mask[:1] + mask
        return mask
//  [python]
def forward(self, x, cond, fea=None):
        if fea is None:
            feat = self.conv_first(x)
        else:
            feat_rgb = self.conv_first(x)
            # feat += torch.sigmoid(self.conv_fea(fea))
            feat = torch.cat((feat_rgb, fea), dim=1)
            feat = self.conv_prefea(feat)
        cond = self.CondNet(cond)
        body_feat = self.body((feat, cond))
        body_feat = self.sftbody(body_feat[0], body_feat[1])
        body_feat = self.conv_body(body_feat)
        body_feat += feat
        # upsample
        if self.scale > 1:
            body_feat = self.lrelu(
                self.conv_up1(
                    t.upsample(body_feat, scale_factor=2, mode='nearest')))
            if self.scale == 4:
                body_feat = self.lrelu(
                    self.conv_up2(
                        t.upsample(
                            body_feat, scale_factor=2, mode='nearest')))
        out = self.conv_last(self.lrelu(self.conv_hr(body_feat)))
        return out
//  [python]
def forward(self, x, cond, fea=None):
        if fea is None:
            feat = self.conv_first(x)
        else:
            feat_rgb = self.conv_first(x)
            # feat += torch.sigmoid(self.conv_fea(fea))
            feat = torch.cat((feat_rgb, fea), dim=1)
            feat = self.conv_prefea(feat)
        cond = self.CondNet(cond)
        body_feat = self.body((feat, cond))
        body_feat = self.sftbody(body_feat[0], body_feat[1])
        body_feat = self.conv_body(body_feat)
        body_feat += feat
        # upsample
        if self.scale > 1:
            body_feat = self.lrelu(
                self.conv_up1(
                    F.interpolate(body_feat, scale_factor=2, mode='nearest')))
            if self.scale == 4:
                body_feat = self.lrelu(
                    self.conv_up2(
                        F.interpolate(
                            body_feat, scale_factor=2, mode='nearest')))
        out = self.conv_last(self.lrelu(self.conv_hr(body_feat)))
        return out
//  [python]
def ascend_txt(model, lats, sideX, sideY, perceptor, percep, gen, tokenizedtxt):
    if gen == 'biggan':
        cutn = 128
        zs = [*lats()]
        out = model(zs[0], zs[1], 1)
        
    elif gen == 'dall-e':
        cutn = 32
        zs = lats()
        out = unmap_pixels(torch.sigmoid(model(zs)[:, :3].float()))

    elif gen == 'stylegan':
        zs = lats.normu.repeat(1,18,1)
        img = model(zs)
        img = t.nn.functional.upsample_bilinear(img, (224, 224))

        img_logits, _text_logits = perceptor(img, tokenizedtxt.cuda())

        return 1/img_logits * 100, img, zs
    
    p_s = []
    for ch in range(cutn):
        # size = int(sideX*torch.zeros(1,).normal_(mean=.8, std=.3).clip(.5, .95))
        size = int(sideX*torch.zeros(1,).normal_(mean=.39, std=.865).clip(.362, .7099))
        offsetx = torch.randint(0, sideX - size, ())
        offsety = torch.randint(0, sideY - size, ())
        apper = out[:, :, offsetx:offsetx + size, offsety:offsety + size]
        apper = pad_augs(apper)
        # apper = kornia_augs(apper, sideX=sideX)
        apper = t.nn.functional.interpolate(apper, (224, 224), mode='nearest')
        p_s.append(apper)
    into = torch.cat(p_s, 0)
    if gen == 'biggan':
        # into = nom((into + 1) / 2)
        up_noise = 0.01649
        into = into + (up_noise)*torch.randn_like(into, requires_grad=True)
        into = nom((into + 1) / 1.8)
    elif gen == 'dall-e':
        into = nom((into + 1) / 2)
    iii = perceptor.encode_image(into)
    llls = zs #lats()

    if gen == 'dall-e':
        return [0, 10*-torch.cosine_similarity(percep, iii).view(-1, 1).T.mean(1), zs]

    lat_l = torch.abs(1 - torch.std(llls[0], dim=1)).mean() + \
            torch.abs(torch.mean(llls[0])).mean() + \
            4*torch.max(torch.square(llls[0]).mean(), lats.thrsh_lat)
    for array in llls[0]:
        mean = torch.mean(array)
        diffs = array - mean
        var = torch.mean(torch.pow(diffs, 2.0))
        std = torch.pow(var, 0.5)
        zscores = diffs / std
        skews = torch.mean(torch.pow(zscores, 3.0))
        kurtoses = torch.mean(torch.pow(zscores, 4.0)) - 3.0
        lat_l = lat_l + torch.abs(kurtoses) / llls[0].shape[0] + torch.abs(skews) / llls[0].shape[0]
    cls_l = ((50*torch.topk(llls[1],largest=False,dim=1,k=999)[0])**2).mean()
    return [lat_l, cls_l, -100*torch.cosine_similarity(percep, iii, dim=-1).mean(), zs]
//  [python]
def ascend_txt(model, lats, sideX, sideY, perceptor, percep, gen, tokenizedtxt):
    if gen == 'biggan':
        cutn = 128
        zs = [*lats()]
        out = model(zs[0], zs[1], 1)
        
    elif gen == 'dall-e':
        cutn = 32
        zs = lats()
        out = unmap_pixels(torch.sigmoid(model(zs)[:, :3].float()))

    elif gen == 'stylegan':
        zs = lats.normu.repeat(1,18,1)
        img = model(zs)
        img = torch.nn.functional.upsample_bilinear(img, (224, 224))

        img_logits, _text_logits = perceptor(img, tokenizedtxt.cuda())

        return 1/img_logits * 100, img, zs
    
    p_s = []
    for ch in range(cutn):
        # size = int(sideX*torch.zeros(1,).normal_(mean=.8, std=.3).clip(.5, .95))
        size = int(sideX*torch.zeros(1,).normal_(mean=.39, std=.865).clip(.362, .7099))
        offsetx = torch.randint(0, sideX - size, ())
        offsety = torch.randint(0, sideY - size, ())
        apper = out[:, :, offsetx:offsetx + size, offsety:offsety + size]
        apper = pad_augs(apper)
        # apper = kornia_augs(apper, sideX=sideX)
        apper = torch.nn.functional.interpolate(apper, (224, 224), mode='nearest')
        p_s.append(apper)
    into = torch.cat(p_s, 0)
    if gen == 'biggan':
        # into = nom((into + 1) / 2)
        up_noise = 0.01649
        into = into + (up_noise)*torch.randn_like(into, requires_grad=True)
        into = nom((into + 1) / 1.8)
    elif gen == 'dall-e':
        into = nom((into + 1) / 2)
    iii = perceptor.encode_image(into)
    llls = zs #lats()

    if gen == 'dall-e':
        return [0, 10*-torch.cosine_similarity(percep, iii).view(-1, 1).T.mean(1), zs]

    lat_l = torch.abs(1 - torch.std(llls[0], dim=1)).mean() + \
            torch.abs(torch.mean(llls[0])).mean() + \
            4*torch.max(torch.square(llls[0]).mean(), lats.thrsh_lat)
    for array in llls[0]:
        mean = torch.mean(array)
        diffs = array - mean
        var = torch.mean(torch.pow(diffs, 2.0))
        std = torch.pow(var, 0.5)
        zscores = diffs / std
        skews = torch.mean(torch.pow(zscores, 3.0))
        kurtoses = torch.mean(torch.pow(zscores, 4.0)) - 3.0
        lat_l = lat_l + torch.abs(kurtoses) / llls[0].shape[0] + torch.abs(skews) / llls[0].shape[0]
    cls_l = ((50*torch.topk(llls[1],largest=False,dim=1,k=999)[0])**2).mean()
    return [lat_l, cls_l, -100*torch.cosine_similarity(percep, iii, dim=-1).mean(), zs]
//  [python]
def _coarse_decoder_pass(self, mel_specs, encoder_outputs, alignments, input_mask):
        """Double Decoder Consistency"""
        T = mel_specs.shape[1]
        if T % self.coarse_decoder.r > 0:
            padding_size = self.coarse_decoder.r - (T % self.coarse_decoder.r)
            mel_specs = torch.nn.functional.pad(mel_specs, (0, 0, 0, padding_size, 0, 0))
        decoder_outputs_backward, alignments_backward, _ = self.coarse_decoder(
            encoder_outputs.detach(), mel_specs, input_mask
        )
        # scale_factor = self.decoder.r_init / self.decoder.r
        alignments_backward = torch.nn.functional.upsample(
            alignments_backward.transpose(1, 2),
            size=alignments.shape[1],
            mode="nearest",
        ).transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward.transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward[:, :T, :]
        return decoder_outputs_backward, alignments_backward
//  [python]
def _coarse_decoder_pass(self, mel_specs, encoder_outputs, alignments, input_mask):
        """Double Decoder Consistency"""
        T = mel_specs.shape[1]
        if T % self.coarse_decoder.r > 0:
            padding_size = self.coarse_decoder.r - (T % self.coarse_decoder.r)
            mel_specs = torch.nn.functional.pad(mel_specs, (0, 0, 0, padding_size, 0, 0))
        decoder_outputs_backward, alignments_backward, _ = self.coarse_decoder(
            encoder_outputs.detach(), mel_specs, input_mask
        )
        # scale_factor = self.decoder.r_init / self.decoder.r
        alignments_backward = torch.nn.functional.interpolate(
            alignments_backward.transpose(1, 2),
            size=alignments.shape[1],
            mode="nearest",
        ).transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward.transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward[:, :T, :]
        return decoder_outputs_backward, alignments_backward
//  [python]
def forward(self, *xs):
        assert len(xs) in (1, 2), 'invalid number of inputs'
        if len(xs) == 1:
            x = self.layer2(xs[0])
        else:
            x = self.layer2(xs[0] + self.layer1(xs[1]))
        x = t.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
        x = self.conv_out(x)
        return x
//  [python]
def forward(self, *xs):
        assert len(xs) in (1, 2), 'invalid number of inputs'
        if len(xs) == 1:
            x = self.layer2(xs[0])
        else:
            x = self.layer2(xs[0] + self.layer1(xs[1]))
        x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
        x = self.conv_out(x)
        return x
//  [python]
def process_mask_upsample(protos, masks_in, bboxes, shape):
    """
    Crop after upsample.
    proto_out: [mask_dim, mask_h, mask_w]
    out_masks: [n, mask_dim], n is number of masks after nms
    bboxes: [n, 4], n is number of masks after nms
    shape:input_image_size, (h, w)

    return: h, w, n
    """

    c, mh, mw = protos.shape  # CHW
    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
    masks = t.nn.functional.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
    masks = crop_mask(masks, bboxes)  # CHW
    return masks.gt_(0.5)
//  [python]
def process_mask_upsample(protos, masks_in, bboxes, shape):
    """
    Crop after upsample.
    proto_out: [mask_dim, mask_h, mask_w]
    out_masks: [n, mask_dim], n is number of masks after nms
    bboxes: [n, 4], n is number of masks after nms
    shape:input_image_size, (h, w)

    return: h, w, n
    """

    c, mh, mw = protos.shape  # CHW
    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
    masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
    masks = crop_mask(masks, bboxes)  # CHW
    return masks.gt_(0.5)
//  [python]
def forward(self, x, scale_factor=1.0):
        if scale_factor == 1.0:
            return x
        else:
            x = t.nn.functional.interpolate(
                x, mode=self.mode, align_corners=False, scale_factor=scale_factor
            )
        return x
//  [python]
def forward(self, x, scale_factor=1.0):
        if scale_factor == 1.0:
            return x
        else:
            x = torch.nn.functional.interpolate(
                x, mode=self.mode, align_corners=False, scale_factor=scale_factor
            )
        return x
//  [python]
def down_to_64(x):
    return t.interpolate(x, (64, 64), mode="area")
//  [python]
def down_to_64(x):
    return F.interpolate(x, (64, 64), mode="area")
//  [python]
def resize_bart_abs_pos_emb(weight: torch.Tensor, max_length: int) -> torch.Tensor:
        """
        Resize position embeddings
        Truncate if sequence length of MBart backbone is greater than given max_length,
        else interpolate to max_length
        """
        if weight.shape[0] > max_length:
            weight = weight[:max_length, ...]
        else:
            weight = (
                t.nn.functional.interpolate(
                    weight.permute(1, 0).unsqueeze(0),
                    size=max_length,
                    mode="linear",
                    align_corners=False,
                )
                .squeeze(0)
                .permute(1, 0)
            )
        return weight