jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x4, x3, x2, x1, x0):
        x = F.interpolate(x4, size=x3.shape[2:], mode='bilinear', align_corners=False)
        x = torch.cat([x, x3], dim=1)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = F.interpolate(x, size=x2.shape[2:], mode='bilinear', align_corners=False)
        x = torch.cat([x, x2], dim=1)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = F.interpolate(x, size=x1.shape[2:], mode='bilinear', align_corners=False)
        x = torch.cat([x, x1], dim=1)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu(x)
        x = F.interpolate(x, size=x0.shape[2:], mode='bilinear', align_corners=False)
        x = torch.cat([x, x0], dim=1)
        x = self.conv4(x)
        return x
//  [python]
def forward(self, x, timesteps, low_res, aug_level=None, **kwargs):
        bs, _, new_height, new_width = x.shape

        align_corners = True
        if self.interpolate_mode == 'nearest':
            align_corners = None

        upsampled = F.interpolate(
            low_res, (new_height, new_width), mode=self.interpolate_mode, align_corners=align_corners
        )

        if aug_level is None:
            aug_steps = (np.random.random(bs)*1000).astype(np.int64)  # uniform [0, 1)
            aug_steps = torch.from_numpy(aug_steps).to(x.device, dtype=torch.long)
        else:
            aug_steps = torch.tensor([int(aug_level * 1000)]).repeat(bs).to(x.device, dtype=torch.long)

        upsampled = self.low_res_diffusion.q_sample(upsampled, aug_steps)
        x = torch.cat([x, upsampled], dim=1)

        aug_emb = self.aug_proj(
            timestep_embedding(aug_steps, self.model_channels, dtype=self.dtype)
        )
        return super().forward(x, timesteps, aug_emb=aug_emb, **kwargs)
//  [python]
def time_warp(self, x):
        """Time warping with torch.nn.functional.interpolate"""
        original_size = x.shape
        window = self.time_warp_window
        if x.dim() == 3:
            x = x.unsqueeze(1)
        time = x.shape[2]
        if time - window <= window:
            return x.view(*original_size)
        c = torch.randint(window, time - window, (1,))[0]
        w = torch.randint(c - window, c + window, (1,))[0] + 1
        left = torch.nn.functional.interpolate(x[:, :, :c], (w, x.shape[3]), mode=self.time_warp_mode, align_corners=True)
        right = torch.nn.functional.interpolate(x[:, :, c:], (time - w, x.shape[3]), mode=self.time_warp_mode, align_corners=True)
        x[:, :, :w] = left
        x[:, :, w:] = right
        return x.view(*original_size)
//  [python]
def forward(self, waveforms, lengths):
        """
        Arguments
        ---------
        waveforms : tensor
            Shape should be `[batch, time]` or `[batch, time, channels]`.
        lengths : tensor
            Shape should be a single dimension, `[batch]`.

        Returns
        -------
        Tensor of shape `[batch, time]` or `[batch, time, channels]`.
        """
        if torch.rand(1) > self.reverb_prob:
            return waveforms.clone()
        channel_added = False
        if len(waveforms.shape) == 2:
            waveforms = waveforms.unsqueeze(-1)
            channel_added = True
        rir_waveform = self._load_rir(waveforms)
        if hasattr(self, 'resampler'):
            rir_waveform = self.resampler(rir_waveform)
        if self.rir_scale_factor != 1:
            rir_waveform = t.interpolate(rir_waveform.transpose(1, -1), scale_factor=self.rir_scale_factor, mode='linear', align_corners=False)
            rir_waveform = rir_waveform.transpose(1, -1)
        rev_waveform = reverberate(waveforms, rir_waveform, rescale_amp='avg')
        if channel_added:
            return rev_waveform.squeeze(-1)
        return rev_waveform
//  [python]
def forward(self, waveforms, lengths):
        """
        Arguments
        ---------
        waveforms : tensor
            Shape should be `[batch, time]` or `[batch, time, channels]`.
        lengths : tensor
            Shape should be a single dimension, `[batch]`.

        Returns
        -------
        Tensor of shape `[batch, time]` or `[batch, time, channels]`.
        """
        if torch.rand(1) > self.reverb_prob:
            return waveforms.clone()
        channel_added = False
        if len(waveforms.shape) == 2:
            waveforms = waveforms.unsqueeze(-1)
            channel_added = True
        rir_waveform = self._load_rir(waveforms)
        if hasattr(self, 'resampler'):
            rir_waveform = self.resampler(rir_waveform)
        if self.rir_scale_factor != 1:
            rir_waveform = F.interpolate(rir_waveform.transpose(1, -1), scale_factor=self.rir_scale_factor, mode='linear', align_corners=False)
            rir_waveform = rir_waveform.transpose(1, -1)
        rev_waveform = reverberate(waveforms, rir_waveform, rescale_amp='avg')
        if channel_added:
            return rev_waveform.squeeze(-1)
        return rev_waveform
//  [python]
def extract_geometry(self, sdf_network, bound_min, bound_max, resolution, threshold, device, occupancy_mask=None,
                         # * 3d feature volume
                         **kwargs
                         ):
        # logging.info('threshold: {}'.format(threshold))

        u = self.extract_fields(bound_min, bound_max, resolution,
                                lambda pts, **kwargs: sdf_network.sdf(pts, **kwargs),
                                # - sdf need to be multiplied by -1
                                device,
                                # * 3d feature volume
                                **kwargs
                                )
        if occupancy_mask is not None:
            dX, dY, dZ = occupancy_mask.shape
            empty_mask = 1 - occupancy_mask
            empty_mask = empty_mask.view(1, 1, dX, dY, dZ)
            # - dilation
            # empty_mask = F.avg_pool3d(empty_mask, kernel_size=7, stride=1, padding=3)
            empty_mask = t.nn.functional.interpolate(empty_mask, [resolution, resolution, resolution], mode='nearest')
            empty_mask = empty_mask.view(resolution, resolution, resolution).cpu().numpy() > 0
            u[empty_mask] = -100
            del empty_mask

        vertices, triangles = mcubes.marching_cubes(u, threshold)
        b_max_np = bound_max.detach().cpu().numpy()
        b_min_np = bound_min.detach().cpu().numpy()

        vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
        return vertices, triangles, u
//  [python]
def extract_geometry(self, sdf_network, bound_min, bound_max, resolution, threshold, device, occupancy_mask=None,
                         # * 3d feature volume
                         **kwargs
                         ):
        # logging.info('threshold: {}'.format(threshold))

        u = self.extract_fields(bound_min, bound_max, resolution,
                                lambda pts, **kwargs: sdf_network.sdf(pts, **kwargs),
                                # - sdf need to be multiplied by -1
                                device,
                                # * 3d feature volume
                                **kwargs
                                )
        if occupancy_mask is not None:
            dX, dY, dZ = occupancy_mask.shape
            empty_mask = 1 - occupancy_mask
            empty_mask = empty_mask.view(1, 1, dX, dY, dZ)
            # - dilation
            # empty_mask = F.avg_pool3d(empty_mask, kernel_size=7, stride=1, padding=3)
            empty_mask = F.interpolate(empty_mask, [resolution, resolution, resolution], mode='nearest')
            empty_mask = empty_mask.view(resolution, resolution, resolution).cpu().numpy() > 0
            u[empty_mask] = -100
            del empty_mask

        vertices, triangles = mcubes.marching_cubes(u, threshold)
        b_max_np = bound_max.detach().cpu().numpy()
        b_min_np = bound_min.detach().cpu().numpy()

        vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
        return vertices, triangles, u
//  [python]
def sample_points(init_normal, gt_norm_mask, sampling_ratio, beta):
    device = init_normal.device
    B, _, H, W = init_normal.shape
    N = int(sampling_ratio * H * W)
    beta = beta

    # uncertainty map
    uncertainty_map = -1 * init_normal[:, 3, :, :]  # B, H, W

    # gt_invalid_mask (B, H, W)
    if gt_norm_mask is not None:
        gt_invalid_mask = t.nn.functional.interpolate(gt_norm_mask.float(), size=[H, W], mode='nearest')
        gt_invalid_mask = gt_invalid_mask[:, 0, :, :] < 0.5
        uncertainty_map[gt_invalid_mask] = -1e4

    # (B, H*W)
    _, idx = uncertainty_map.view(B, -1).sort(1, descending=True)

    # importance sampling
    if int(beta * N) > 0:
        importance = idx[:, :int(beta * N)]    # B, beta*N

        # remaining
        remaining = idx[:, int(beta * N):]     # B, H*W - beta*N

        # coverage
        num_coverage = N - int(beta * N)

        if num_coverage <= 0:
            samples = importance
        else:
            coverage_list = []
            for i in range(B):
                idx_c = torch.randperm(remaining.size()[1])    # shuffles "H*W - beta*N"
                coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))     # 1, N-beta*N
            coverage = torch.cat(coverage_list, dim=0)                                      # B, N-beta*N
            samples = torch.cat((importance, coverage), dim=1)                              # B, N

    else:
        # remaining
        remaining = idx[:, :]  # B, H*W

        # coverage
        num_coverage = N

        coverage_list = []
        for i in range(B):
            idx_c = torch.randperm(remaining.size()[1])  # shuffles "H*W - beta*N"
            coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))  # 1, N-beta*N
        coverage = torch.cat(coverage_list, dim=0)  # B, N-beta*N
        samples = coverage

    # point coordinates
    rows_int = samples // W         # 0 for first row, H-1 for last row
    rows_float = rows_int / float(H-1)         # 0 to 1.0
    rows_float = (rows_float * 2.0) - 1.0       # -1.0 to 1.0

    cols_int = samples % W          # 0 for first column, W-1 for last column
    cols_float = cols_int / float(W-1)         # 0 to 1.0
    cols_float = (cols_float * 2.0) - 1.0       # -1.0 to 1.0

    point_coords = torch.zeros(B, 1, N, 2)
    point_coords[:, 0, :, 0] = cols_float             # x coord
    point_coords[:, 0, :, 1] = rows_float             # y coord
    point_coords = point_coords.to(device)
    return point_coords, rows_int, cols_int
//  [python]
def sample_points(init_normal, gt_norm_mask, sampling_ratio, beta):
    device = init_normal.device
    B, _, H, W = init_normal.shape
    N = int(sampling_ratio * H * W)
    beta = beta

    # uncertainty map
    uncertainty_map = -1 * init_normal[:, 3, :, :]  # B, H, W

    # gt_invalid_mask (B, H, W)
    if gt_norm_mask is not None:
        gt_invalid_mask = F.interpolate(gt_norm_mask.float(), size=[H, W], mode='nearest')
        gt_invalid_mask = gt_invalid_mask[:, 0, :, :] < 0.5
        uncertainty_map[gt_invalid_mask] = -1e4

    # (B, H*W)
    _, idx = uncertainty_map.view(B, -1).sort(1, descending=True)

    # importance sampling
    if int(beta * N) > 0:
        importance = idx[:, :int(beta * N)]    # B, beta*N

        # remaining
        remaining = idx[:, int(beta * N):]     # B, H*W - beta*N

        # coverage
        num_coverage = N - int(beta * N)

        if num_coverage <= 0:
            samples = importance
        else:
            coverage_list = []
            for i in range(B):
                idx_c = torch.randperm(remaining.size()[1])    # shuffles "H*W - beta*N"
                coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))     # 1, N-beta*N
            coverage = torch.cat(coverage_list, dim=0)                                      # B, N-beta*N
            samples = torch.cat((importance, coverage), dim=1)                              # B, N

    else:
        # remaining
        remaining = idx[:, :]  # B, H*W

        # coverage
        num_coverage = N

        coverage_list = []
        for i in range(B):
            idx_c = torch.randperm(remaining.size()[1])  # shuffles "H*W - beta*N"
            coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))  # 1, N-beta*N
        coverage = torch.cat(coverage_list, dim=0)  # B, N-beta*N
        samples = coverage

    # point coordinates
    rows_int = samples // W         # 0 for first row, H-1 for last row
    rows_float = rows_int / float(H-1)         # 0 to 1.0
    rows_float = (rows_float * 2.0) - 1.0       # -1.0 to 1.0

    cols_int = samples % W          # 0 for first column, W-1 for last column
    cols_float = cols_int / float(W-1)         # 0 to 1.0
    cols_float = (cols_float * 2.0) - 1.0       # -1.0 to 1.0

    point_coords = torch.zeros(B, 1, N, 2)
    point_coords[:, 0, :, 0] = cols_float             # x coord
    point_coords[:, 0, :, 1] = rows_float             # y coord
    point_coords = point_coords.to(device)
    return point_coords, rows_int, cols_int
//  [python]
def encode_decode(self, inputs: Tensor,
                      batch_img_metas: List[dict]) -> Tensor:
        """Encode the name of classes with text_encoder and encode images with
        image_encoder.

        Then decode the class embedding and visual feature into a semantic
        segmentation map of the same size as input.
        """
        classifier_embeds = self.text_encoder()
        clip_inputs = inputs
        if self.asymetric_input:
            clip_inputs = t.interpolate(
                inputs, scale_factor=self.encoder_resolution, mode='bilinear')
        x = self.image_encoder(clip_inputs)
        seg_logits = self.decode_head.predict([inputs, x, classifier_embeds],
                                              batch_img_metas, self.test_cfg)

        return seg_logits
//  [python]
def encode_decode(self, inputs: Tensor,
                      batch_img_metas: List[dict]) -> Tensor:
        """Encode the name of classes with text_encoder and encode images with
        image_encoder.

        Then decode the class embedding and visual feature into a semantic
        segmentation map of the same size as input.
        """
        classifier_embeds = self.text_encoder()
        clip_inputs = inputs
        if self.asymetric_input:
            clip_inputs = F.interpolate(
                inputs, scale_factor=self.encoder_resolution, mode='bilinear')
        x = self.image_encoder(clip_inputs)
        seg_logits = self.decode_head.predict([inputs, x, classifier_embeds],
                                              batch_img_metas, self.test_cfg)

        return seg_logits
//  [python]
def loss(self, inputs: Tensor, data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            inputs (Tensor): Input images.
            data_samples (list[:obj:`SegDataSample`]): The seg data samples.
                It usually includes information such as `metainfo` and
                `gt_sem_seg`.

        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        classifier_embeds = self.text_encoder()
        clip_inputs = inputs
        if self.asymetric_input:
            clip_inputs = t.interpolate(
                inputs, scale_factor=self.encoder_resolution, mode='bilinear')
        x = self.image_encoder(clip_inputs)

        losses = dict()

        loss_decode = self._decode_head_forward_train(
            [inputs, x, classifier_embeds], data_samples)
        losses.update(loss_decode)

        return losses
//  [python]
def loss(self, inputs: Tensor, data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            inputs (Tensor): Input images.
            data_samples (list[:obj:`SegDataSample`]): The seg data samples.
                It usually includes information such as `metainfo` and
                `gt_sem_seg`.

        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        classifier_embeds = self.text_encoder()
        clip_inputs = inputs
        if self.asymetric_input:
            clip_inputs = F.interpolate(
                inputs, scale_factor=self.encoder_resolution, mode='bilinear')
        x = self.image_encoder(clip_inputs)

        losses = dict()

        loss_decode = self._decode_head_forward_train(
            [inputs, x, classifier_embeds], data_samples)
        losses.update(loss_decode)

        return losses
//  [python]
def evaluate(self, outs):
    """Plot a frame along with its reconstruction."""

    def _remove_ctx_frames(frame):
      s, h, w, c = frame.shape
      seq_len = s // self.num_ctx_frames
      frame = frame.reshape(seq_len, self.num_ctx_frames, h, w, c)
      return frame[:, -1]

    frames = [o.frames for o in outs]
    recons = [o.reconstruction for o in outs]

    r_idx = np.random.randint(0, len(frames))
    frame = _remove_ctx_frames(frames[r_idx])
    recon = _remove_ctx_frames(recons[r_idx])

    # Select which frames we want to plot from the sequence.
    frame_idxs = np.random.choice(
        np.arange(frame.shape[0]), size=self.num_frames, replace=False)
    frame = frame[frame_idxs]
    recon = recon[frame_idxs]

    # Downsample the frame.
    _, _, sh, _ = recon.shape
    _, _, h, _ = frame.shape
    scale_factor = sh / h
    frame_ds = t.nn.functional.interpolate(
        torch.from_numpy(frame).permute(0, 3, 1, 2),
        mode='bilinear',
        scale_factor=scale_factor,
        recompute_scale_factor=False,
        align_corners=True).permute(0, 2, 3, 1).numpy()

    # Clip reconstruction between 0 and 1.
    recon = np.clip(recon, 0.0, 1.0)

    imgs = np.concatenate([frame_ds, recon], axis=0)
    img = make_grid(torch.from_numpy(imgs).permute(0, 3, 1, 2), nrow=2)

    return EvaluatorOutput(image=img.permute(1, 2, 0).numpy())
//  [python]
def evaluate(self, outs):
    """Plot a frame along with its reconstruction."""

    def _remove_ctx_frames(frame):
      s, h, w, c = frame.shape
      seq_len = s // self.num_ctx_frames
      frame = frame.reshape(seq_len, self.num_ctx_frames, h, w, c)
      return frame[:, -1]

    frames = [o.frames for o in outs]
    recons = [o.reconstruction for o in outs]

    r_idx = np.random.randint(0, len(frames))
    frame = _remove_ctx_frames(frames[r_idx])
    recon = _remove_ctx_frames(recons[r_idx])

    # Select which frames we want to plot from the sequence.
    frame_idxs = np.random.choice(
        np.arange(frame.shape[0]), size=self.num_frames, replace=False)
    frame = frame[frame_idxs]
    recon = recon[frame_idxs]

    # Downsample the frame.
    _, _, sh, _ = recon.shape
    _, _, h, _ = frame.shape
    scale_factor = sh / h
    frame_ds = F.interpolate(
        torch.from_numpy(frame).permute(0, 3, 1, 2),
        mode='bilinear',
        scale_factor=scale_factor,
        recompute_scale_factor=False,
        align_corners=True).permute(0, 2, 3, 1).numpy()

    # Clip reconstruction between 0 and 1.
    recon = np.clip(recon, 0.0, 1.0)

    imgs = np.concatenate([frame_ds, recon], axis=0)
    img = make_grid(torch.from_numpy(imgs).permute(0, 3, 1, 2), nrow=2)

    return EvaluatorOutput(image=img.permute(1, 2, 0).numpy())
//  [python]
def test_interpolate_without_scales_nearest_5d(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 12, 16)
        output_shape = (16, 24, 32)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'nearest'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(), outputs['output'])
//  [python]
def test_interpolate_without_scales_nearest_5d(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 12, 16)
        output_shape = (16, 24, 32)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'nearest'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(), outputs['output'])
//  [python]
def test_interpolate_without_scales_bilinear_4d_disable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 12)
        output_shape = (16, 24)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bilinear'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_without_scales_bilinear_4d_disable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 12)
        output_shape = (16, 24)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bilinear'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_with_scale_3d_nearest_exact(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 16)
        scales_factor = (2, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'nearest-exact'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(), outputs['output'])
//  [python]
def test_interpolate_with_scale_3d_nearest_exact(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 16)
        scales_factor = (2, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'nearest-exact'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(), outputs['output'])
//  [python]
def test_interpolate_without_scales_bicubic_4d_enable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 12)
        output_shape = (16, 24)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bicubic'
        # construct trt network
        align_corners_flag = True

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.nn.functional.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-3)
//  [python]
def test_interpolate_without_scales_bicubic_4d_enable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 12)
        output_shape = (16, 24)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bicubic'
        # construct trt network
        align_corners_flag = True

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-3)
//  [python]
def test_interpolate_without_scales_bilinear_4d_enable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 12)
        output_shape = (16, 24)

        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bilinear'
        # construct trt network
        align_corners_flag = True

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                size=output_shape,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              size=output_shape,
                                              mode=mode)

        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_with_scale_4d_bicubic(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 12)
        scales_factor = (2.5, 2)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bicubic'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_with_scale_4d_bicubic(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 4, 8, 12)
        scales_factor = (2.5, 2)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bicubic'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def resize_mask(mask, shape):
    return t.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
//  [python]
def resize_mask(mask, shape):
    return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        if self.use_conv_transpose:
            return self.conv(x)

        x = F.upsample(x, scale_factor=2.0, mode="nearest")

        if self.use_conv:
            x = self.conv(x)

        return x
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        if self.use_conv_transpose:
            return self.conv(x)

        x = F.interpolate(x, scale_factor=2.0, mode="nearest")

        if self.use_conv:
            x = self.conv(x)

        return x
//  [python]
def resize_image_to(
    image,
    target_image_size,
    clamp_range = None,
    mode = 'nearest'
):
    orig_image_size = image.shape[-1]

    if orig_image_size == target_image_size:
        return image

    out = t.interpolate(image, target_image_size, mode = mode)

    if exists(clamp_range):
        out = out.clamp(*clamp_range)

    return out
//  [python]
def resize_image_to(
    image,
    target_image_size,
    clamp_range = None,
    mode = 'nearest'
):
    orig_image_size = image.shape[-1]

    if orig_image_size == target_image_size:
        return image

    out = F.interpolate(image, target_image_size, mode = mode)

    if exists(clamp_range):
        out = out.clamp(*clamp_range)

    return out
//  [python]
def test_interpolate_with_scale_4d_bilinear(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 32)
        scales_factor = (2.5, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bilinear'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              align_corners=align_corners_flag,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_with_scale_4d_bilinear(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 32)
        scales_factor = (2.5, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'bilinear'
        # construct trt network
        align_corners_flag = False

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              align_corners=align_corners_flag,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def inference(self, c, g=None):
        """
        Args:
            x (Tensor): conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """
        # c = c.to(self.conv_pre.weight.device)
        # c = torch.nn.functional.pad(c, (self.inference_padding, self.inference_padding), "replicate")
        up_1 = torch.nn.functional.upsample(
                c.transpose(1,2),
                scale_factor=[1024 / 256],
                mode="linear",
            )
        up_2 = torch.nn.functional.upsample(
            up_1,
            scale_factor=[24000 / 22050],
            mode="linear",
        )
        g = g.unsqueeze(0)
        return self.forward(up_2.to(self.device), g.transpose(1,2))
//  [python]
def inference(self, c, g=None):
        """
        Args:
            x (Tensor): conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """
        # c = c.to(self.conv_pre.weight.device)
        # c = torch.nn.functional.pad(c, (self.inference_padding, self.inference_padding), "replicate")
        up_1 = torch.nn.functional.interpolate(
                c.transpose(1,2),
                scale_factor=[1024 / 256],
                mode="linear",
            )
        up_2 = torch.nn.functional.interpolate(
            up_1,
            scale_factor=[24000 / 22050],
            mode="linear",
        )
        g = g.unsqueeze(0)
        return self.forward(up_2.to(self.device), g.transpose(1,2))
//  [python]
def test_interpolate_with_scale_5d_trilinear_enable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 16, 32)
        scales_factor = (2.5, 2, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'trilinear'
        # construct trt network
        align_corners_flag = True

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = t.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              align_corners=align_corners_flag,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def test_interpolate_with_scale_5d_trilinear_enable_align_corner(self):
        # test data
        dtype = 'float32'
        input_shape = (1, 1, 8, 16, 32)
        scales_factor = (2.5, 2, 4)
        input_data = torch.rand(
            input_shape, dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype))
        mode = 'trilinear'
        # construct trt network
        align_corners_flag = True

        builder = tensorrt_llm.Builder()
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            input = Tensor(name='input',
                           shape=input_shape,
                           dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.interpolate(
                input=input,
                scale_factor=scales_factor,
                mode=mode,
                align_corners=align_corners_flag,
            ).trt_tensor
            output.name = 'output'
            network.mark_output(output)

        # trt run
        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))
        with TrtRunner(build_engine) as runner:
            outputs = runner.infer(feed_dict={'input': input_data.numpy()})

        ref = torch.nn.functional.interpolate(input_data,
                                              scale_factor=scales_factor,
                                              align_corners=align_corners_flag,
                                              mode=mode)
        np.testing.assert_allclose(ref.cpu().numpy(),
                                   outputs['output'],
                                   atol=1e-5)
//  [python]
def forward(self, x):
        return F.interpolate(x, size=self.size, scale_factor=self.scale_factor,mode=self.mode, align_corners=self.align_corners)
//  [python]
def forward(self, x):
        out = F.upsample(x, scale_factor=(1, 2, 2), mode='trilinear')
        # out = F.interpolate(x, scale_factor=(1, 2, 2))
        out = self.conv(out)
        out = self.norm(out)
        out = F.relu(out)
        return out
//  [python]
def forward(self, x):
        # out = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear')
        out = F.interpolate(x, scale_factor=(1, 2, 2))
        out = self.conv(out)
        out = self.norm(out)
        out = F.relu(out)
        return out
//  [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 t.interpolate(input, size, scale_factor, mode, align_corners)