jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_where_broadcast(input_shapes):
    torch.manual_seed(42)

    cond_shape = input_shapes[0]
    x_shape = input_shapes[1]
    y_shape = input_shapes[2]

    class Model(torch.nn.Module):
        def forward(self, cond, x, y):
            return torch.where(cond, x, y)

    cond = torch.empty(cond_shape).bernoulli_().to(torch.bool)
    x = torch.randn(x_shape)
    y = torch.randn(y_shape)

    cpu_mod = Model()
    ipu_mod = poptorch.inferenceModel(cpu_mod)

    torch.testing.assert_close(actual=ipu_mod(cond, x, y),
                               expected=cpu_mod(cond, x, y))
//  [python]
def test_with_randomized_data_getter(self) -> None:
        num_samples = 7
        num_tasks = 5
        for _ in range(100):
            input, target = rd.get_rand_data_binary(
                num_updates=num_samples, num_tasks=num_tasks, batch_size=1
            )
            target = target.reshape(num_tasks, num_samples)
            input = input.reshape(num_tasks, num_samples)

            for k in [1, num_samples // 2, num_samples, num_samples * 2, None]:
                actual_prec = retrieval_precision(
                    input,
                    target,
                    k=k,
                    num_tasks=num_tasks,
                )
                expected_prec = torch.stack(
                    [self.rp_1D(inp, tar, k=k) for inp, tar in zip(input, target)]
                )
                t.assert_close(actual_prec, expected_prec)
//  [python]
def test_with_randomized_data_getter(self) -> None:
        num_samples = 7
        num_tasks = 5
        for _ in range(100):
            input, target = rd.get_rand_data_binary(
                num_updates=num_samples, num_tasks=num_tasks, batch_size=1
            )
            target = target.reshape(num_tasks, num_samples)
            input = input.reshape(num_tasks, num_samples)

            for k in [1, num_samples // 2, num_samples, num_samples * 2, None]:
                actual_prec = retrieval_precision(
                    input,
                    target,
                    k=k,
                    num_tasks=num_tasks,
                )
                expected_prec = torch.stack(
                    [self.rp_1D(inp, tar, k=k) for inp, tar in zip(input, target)]
                )
                torch.testing.assert_close(actual_prec, expected_prec)
//  [python]
def test_logn_attention(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        eps = 1e-6
        for batch_size in range(32, 4096, 512):
            for hidden_units in range(1024, 4096, 1024):

                self.gamm = torch.rand(hidden_units).cuda()
                self.T5LayerNormOp = torch.classes.unittest.T5LayerNormOp(eps)
                self.T5LayerNorm = T5LayerNorm(hidden_units, eps)
                self.T5LayerNorm.weight = torch.nn.Parameter(self.gamm)

                hidden_states = torch.rand(batch_size, hidden_units).cuda()
                result = self.T5LayerNorm(hidden_states)
                test = self.T5LayerNormOp.forward(hidden_states, self.gamm)
                torch.testing.assert_allclose(test, result)
//  [python]
def test_logn_attention(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        eps = 1e-6
        for batch_size in range(32, 4096, 512):
            for hidden_units in range(1024, 4096, 1024):

                self.gamm = torch.rand(hidden_units).cuda()
                self.T5LayerNormOp = torch.classes.unittest.T5LayerNormOp(eps)
                self.T5LayerNorm = T5LayerNorm(hidden_units, eps)
                self.T5LayerNorm.weight = torch.nn.Parameter(self.gamm)

                hidden_states = torch.rand(batch_size, hidden_units).cuda()
                result = self.T5LayerNorm(hidden_states)
                test = self.T5LayerNormOp.forward(hidden_states, self.gamm)
                torch.testing.assert_close(test, result)
//  [python]
def _validate(
        cls,
        actual_output: torch.Tensor,
        expected_output: torch.Tensor,
        input: torch.Tensor,
        world_size: int,
    ) -> None:
        t.assert_allclose(
            actual_output.cpu().detach(), expected_output.cpu().detach()
        )
        t.assert_allclose(
            # pyre-fixme[16]: Optional type has no attribute `cpu`.
            input.grad.cpu().detach(),
            torch.ones(input.size()),
        )
//  [python]
def _validate(
        cls,
        actual_output: torch.Tensor,
        expected_output: torch.Tensor,
        input: torch.Tensor,
        world_size: int,
    ) -> None:
        torch.testing.assert_close(
            actual_output.cpu().detach(), expected_output.cpu().detach()
        )
        torch.testing.assert_close(
            # pyre-fixme[16]: Optional type has no attribute `cpu`.
            input.grad.cpu().detach(),
            torch.ones(input.size()),
        )
//  [python]
def test_with_example_cases_one_task(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        t.assert_close(
            retrieval_precision(
                input,
                target,
            ),
            torch.tensor(4 / 7),
        )
        t.assert_close(
            retrieval_precision(
                input,
                target,
                k=2,
            ),
            torch.tensor(1 / 2),
        )
        t.assert_close(
            retrieval_precision(
                input,
                target,
                k=4,
            ),
            torch.tensor(1 / 2),
        )
        t.assert_close(
            retrieval_precision(
                input,
                target,
                k=10,
            ),
            torch.tensor(2 / 5),
        )
        t.assert_close(
            retrieval_precision(
                input,
                target,
                k=10,
                limit_k_to_size=True,
            ),
            torch.tensor(4 / 7),
        )
//  [python]
def test_with_example_cases_one_task(self) -> None:
        input = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        torch.testing.assert_close(
            retrieval_precision(
                input,
                target,
            ),
            torch.tensor(4 / 7),
        )
        torch.testing.assert_close(
            retrieval_precision(
                input,
                target,
                k=2,
            ),
            torch.tensor(1 / 2),
        )
        torch.testing.assert_close(
            retrieval_precision(
                input,
                target,
                k=4,
            ),
            torch.tensor(1 / 2),
        )
        torch.testing.assert_close(
            retrieval_precision(
                input,
                target,
                k=10,
            ),
            torch.tensor(2 / 5),
        )
        torch.testing.assert_close(
            retrieval_precision(
                input,
                target,
                k=10,
                limit_k_to_size=True,
            ),
            torch.tensor(4 / 7),
        )
//  [python]
def test_with_example_cases_multi_tasks_k_None(self) -> None:
        k = None
        input = torch.tensor([[0.1, 0.2, 0.3]]).repeat(8, 1)
        target = torch.tensor(
            [
                [0, 0, 0],
                [0, 0, 1],
                [0, 1, 0],
                [0, 1, 1],
                [1, 0, 0],
                [1, 0, 1],
                [1, 1, 0],
                [1, 1, 1],
            ]
        )
        expected_rp = torch.tensor([0.0, 1 / 3, 1 / 3, 2 / 3, 1 / 3, 2 / 3, 2 / 3, 1.0])
        t.assert_close(
            retrieval_precision(input, target, k=k, num_tasks=8), expected_rp
        )
        t.assert_close(
            retrieval_precision(input, target, k=10, limit_k_to_size=True, num_tasks=8),
            expected_rp,
        )
//  [python]
def test_with_example_cases_multi_tasks_k_None(self) -> None:
        k = None
        input = torch.tensor([[0.1, 0.2, 0.3]]).repeat(8, 1)
        target = torch.tensor(
            [
                [0, 0, 0],
                [0, 0, 1],
                [0, 1, 0],
                [0, 1, 1],
                [1, 0, 0],
                [1, 0, 1],
                [1, 1, 0],
                [1, 1, 1],
            ]
        )
        expected_rp = torch.tensor([0.0, 1 / 3, 1 / 3, 2 / 3, 1 / 3, 2 / 3, 2 / 3, 1.0])
        torch.testing.assert_close(
            retrieval_precision(input, target, k=k, num_tasks=8), expected_rp
        )
        torch.testing.assert_close(
            retrieval_precision(input, target, k=10, limit_k_to_size=True, num_tasks=8),
            expected_rp,
        )
//  [python]
def _validate(
        cls,
        actual_output_awaitable: Union[KJTAllToAllSplitsAwaitable, KeyedJaggedTensor],
        expected_output_awaitable: Union[KJTAllToAllSplitsAwaitable, KeyedJaggedTensor],
    ) -> None:
        actual_output = (
            actual_output_awaitable
            if isinstance(actual_output_awaitable, KeyedJaggedTensor)
            else actual_output_awaitable.wait().wait()
        )
        expected_output = (
            expected_output_awaitable
            if isinstance(expected_output_awaitable, KeyedJaggedTensor)
            else expected_output_awaitable.wait().wait()
        )
        t.assert_close(
            actual_output.values().cpu(),
            expected_output.values().cpu(),
        )
        t.assert_close(
            (
                actual_output.weights().cpu()
                if actual_output.weights_or_none() is not None
                else []
            ),
            (
                expected_output.weights().cpu()
                if expected_output.weights_or_none() is not None
                else []
            ),
        )
        t.assert_close(
            actual_output.lengths().cpu(),
            expected_output.lengths().cpu(),
        )
        assert actual_output.keys() == expected_output.keys()
//  [python]
def _validate(
        cls,
        actual_output_awaitable: Union[KJTAllToAllSplitsAwaitable, KeyedJaggedTensor],
        expected_output_awaitable: Union[KJTAllToAllSplitsAwaitable, KeyedJaggedTensor],
    ) -> None:
        actual_output = (
            actual_output_awaitable
            if isinstance(actual_output_awaitable, KeyedJaggedTensor)
            else actual_output_awaitable.wait().wait()
        )
        expected_output = (
            expected_output_awaitable
            if isinstance(expected_output_awaitable, KeyedJaggedTensor)
            else expected_output_awaitable.wait().wait()
        )
        torch.testing.assert_close(
            actual_output.values().cpu(),
            expected_output.values().cpu(),
        )
        torch.testing.assert_close(
            (
                actual_output.weights().cpu()
                if actual_output.weights_or_none() is not None
                else []
            ),
            (
                expected_output.weights().cpu()
                if expected_output.weights_or_none() is not None
                else []
            ),
        )
        torch.testing.assert_close(
            actual_output.lengths().cpu(),
            expected_output.lengths().cpu(),
        )
        assert actual_output.keys() == expected_output.keys()
//  [python]
def _run_test_dist(
        cls,
        rank: int,
        world_size: int,
        input: torch.Tensor,
        input_splits: List[int],
        expected_output: torch.Tensor,
        qcomms_config: Optional[QCommsConfig] = None,
    ) -> None:
        dist.init_process_group(rank=rank, world_size=2, backend="nccl")
        pg = dist.group.WORLD
        input = input.cuda(rank)
        input.requires_grad = True

        codecs = get_qcomm_codecs(qcomms_config)

        rs = PooledEmbeddingsReduceScatter(
            # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
            #  `Optional[_distributed_c10d.ProcessGroup]`.
            pg,
            codecs=codecs,
        ).cuda(rank)
        actual_output = rs(input, input_splits=input_splits).wait()
        s = torch.sum(actual_output)
        s.backward()

        atol, rtol = None, None
        if qcomms_config is not None:
            atol, rtol = 0.003, 0.004
        torch.testing.assert_allclose(
            actual_output.cpu().detach(),
            expected_output.cpu().detach(),
            rtol=rtol,
            atol=atol,
        )
        if qcomms_config is None:
            torch.testing.assert_allclose(
                # pyre-ignore
                input.grad.cpu().detach(),
                torch.ones(input.size()).div_(world_size),
            )
//  [python]
def _run_test_dist(
        cls,
        rank: int,
        world_size: int,
        input: torch.Tensor,
        input_splits: List[int],
        expected_output: torch.Tensor,
        qcomms_config: Optional[QCommsConfig] = None,
    ) -> None:
        dist.init_process_group(rank=rank, world_size=2, backend="nccl")
        pg = dist.group.WORLD
        input = input.cuda(rank)
        input.requires_grad = True

        codecs = get_qcomm_codecs(qcomms_config)

        rs = PooledEmbeddingsReduceScatter(
            # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
            #  `Optional[_distributed_c10d.ProcessGroup]`.
            pg,
            codecs=codecs,
        ).cuda(rank)
        actual_output = rs(input, input_splits=input_splits).wait()
        s = torch.sum(actual_output)
        s.backward()

        atol, rtol = None, None
        if qcomms_config is not None:
            atol, rtol = 0.003, 0.004
        torch.testing.assert_close(
            actual_output.cpu().detach(),
            expected_output.cpu().detach(),
            rtol=rtol,
            atol=atol,
        )
        if qcomms_config is None:
            torch.testing.assert_close(
                # pyre-ignore
                input.grad.cpu().detach(),
                torch.ones(input.size()).div_(world_size),
            )
//  [python]
def _run_test_dist(
        cls,
        rank: int,
        world_size: int,
        input: torch.Tensor,
        expected_output: torch.Tensor,
        qcomms_config: Optional[QCommsConfig] = None,
    ) -> None:
        dist.init_process_group(rank=rank, world_size=2, backend="nccl")
        pg = dist.group.WORLD
        input = input.cuda(rank)
        input.requires_grad = True

        codecs = get_qcomm_codecs(qcomms_config)

        rs = PooledEmbeddingsReduceScatter(
            # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
            #  `Optional[_distributed_c10d.ProcessGroup]`.
            pg,
            codecs=codecs,
        ).cuda(rank)
        actual_output = rs(input).wait()
        s = torch.sum(actual_output)
        s.backward()

        atol, rtol = None, None
        if qcomms_config is not None:
            atol, rtol = 0.003, 0.004
        torch.testing.assert_allclose(
            actual_output.cpu().detach(),
            expected_output.cpu().detach(),
            rtol=rtol,
            atol=atol,
        )
        if qcomms_config is None:
            torch.testing.assert_allclose(
                # pyre-ignore
                input.grad.cpu().detach(),
                torch.ones(input.size()).div_(world_size),
            )
//  [python]
def _run_test_dist(
        cls,
        rank: int,
        world_size: int,
        input: torch.Tensor,
        expected_output: torch.Tensor,
        qcomms_config: Optional[QCommsConfig] = None,
    ) -> None:
        dist.init_process_group(rank=rank, world_size=2, backend="nccl")
        pg = dist.group.WORLD
        input = input.cuda(rank)
        input.requires_grad = True

        codecs = get_qcomm_codecs(qcomms_config)

        rs = PooledEmbeddingsReduceScatter(
            # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
            #  `Optional[_distributed_c10d.ProcessGroup]`.
            pg,
            codecs=codecs,
        ).cuda(rank)
        actual_output = rs(input).wait()
        s = torch.sum(actual_output)
        s.backward()

        atol, rtol = None, None
        if qcomms_config is not None:
            atol, rtol = 0.003, 0.004
        torch.testing.assert_close(
            actual_output.cpu().detach(),
            expected_output.cpu().detach(),
            rtol=rtol,
            atol=atol,
        )
        if qcomms_config is None:
            torch.testing.assert_close(
                # pyre-ignore
                input.grad.cpu().detach(),
                torch.ones(input.size()).div_(world_size),
            )
//  [python]
def assert_cache(
    tensor_a: torch.Tensor, tensor_b: torch.Tensor, stochastic_rounding: bool
) -> None:
    if stochastic_rounding:
        # Stochastic rounding randomly alters the mantissa bits during the
        # FP32->FP16 conversion in TBE backward, resulting in non-deterministic
        # results. The relative difference between the results from two
        # different runs must be <= 1.0e-2. We set absolute tolerance to 1.0e-2
        # based on the initial embedding weights.
        torch.testing.assert_allclose(
            tensor_a.float(), tensor_b.float(), atol=1.0e-2, rtol=1.0e-2
        )
    else:
        assert_torch_equal(tensor_a, tensor_b)
//  [python]
def assert_cache(
    tensor_a: torch.Tensor, tensor_b: torch.Tensor, stochastic_rounding: bool
) -> None:
    if stochastic_rounding:
        # Stochastic rounding randomly alters the mantissa bits during the
        # FP32->FP16 conversion in TBE backward, resulting in non-deterministic
        # results. The relative difference between the results from two
        # different runs must be <= 1.0e-2. We set absolute tolerance to 1.0e-2
        # based on the initial embedding weights.
        torch.testing.assert_close(
            tensor_a.float(), tensor_b.float(), atol=1.0e-2, rtol=1.0e-2
        )
    else:
        assert_torch_equal(tensor_a, tensor_b)
//  [python]
def test_isin_2dim(dim):
    key = "tensor1"
    input = TensorDict(
        {
            "tensor1": torch.ones(4, 4),
            "tensor2": torch.ones(4, 4),
        },
        batch_size=[4, 4],
    )
    td_ref = TensorDict(
        {
            "tensor1": torch.ones(4, 4),
            "tensor2": torch.ones(4, 4),
        },
        batch_size=[4, 4],
    )

    positive_dim = dim if dim >= 0 else dim + 2
    input[key][(slice(None),) * positive_dim + (0,)] = 2
    in_ref = isin(input, td_ref, key, dim)
    expected_in_ref = torch.tensor([False, True, True, True])
    torch.testing.assert_allclose(in_ref, expected_in_ref)
//  [python]
def test_isin_2dim(dim):
    key = "tensor1"
    input = TensorDict(
        {
            "tensor1": torch.ones(4, 4),
            "tensor2": torch.ones(4, 4),
        },
        batch_size=[4, 4],
    )
    td_ref = TensorDict(
        {
            "tensor1": torch.ones(4, 4),
            "tensor2": torch.ones(4, 4),
        },
        batch_size=[4, 4],
    )

    positive_dim = dim if dim >= 0 else dim + 2
    input[key][(slice(None),) * positive_dim + (0,)] = 2
    in_ref = isin(input, td_ref, key, dim)
    expected_in_ref = torch.tensor([False, True, True, True])
    torch.testing.assert_close(in_ref, expected_in_ref)
//  [python]
def test_convert(tmp_path):
    model = resnet18().cuda().eval()

    trt_model = module2trt(
        model,
        args=[torch.rand(1, 3, 32, 32).cuda()],
    )

    model_path = tmp_path / 'tmp.pth'
    torch.save(trt_model.state_dict(), model_path)
    assert model_path.exists()

    trt_model = TRTModule()
    trt_model.load_state_dict(torch.load(model_path))

    x = torch.rand(1, 3, 32, 32).cuda()
    with torch.no_grad():
        y = model(x)
        y_trt = trt_model(x)

    print(y)
    print(y_trt)
    torch.testing.assert_allclose(y, y_trt)
//  [python]
def test_convert(tmp_path):
    model = resnet18().cuda().eval()

    trt_model = module2trt(
        model,
        args=[torch.rand(1, 3, 32, 32).cuda()],
    )

    model_path = tmp_path / 'tmp.pth'
    torch.save(trt_model.state_dict(), model_path)
    assert model_path.exists()

    trt_model = TRTModule()
    trt_model.load_state_dict(torch.load(model_path))

    x = torch.rand(1, 3, 32, 32).cuda()
    with torch.no_grad():
        y = model(x)
        y_trt = trt_model(x)

    print(y)
    print(y_trt)
    torch.testing.assert_close(y, y_trt)
//  [python]
def test_gelu(self, dtype, strongly_typed):
        if dtype == 'bfloat16':
            self.skip_bf16_before_ampere()

        torch_dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype)
        x_shape = (12, 12, 96, 96)
        x_data = torch.rand(x_shape, dtype=torch_dtype)

        # construct trt network
        builder = tensorrt_llm.Builder()
        builder.strongly_typed = strongly_typed
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            x = Tensor(name='x',
                       shape=x_shape,
                       dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.gelu(x).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={'x': x_data})
        out = outputs['output'].to(torch_dtype)

        # Reference
        ref = self.gelu(x_data, dtype)

        if dtype == 'bfloat16':
            atol, rtol = 1e-5, 2e-2
        else:
            atol, rtol = 1e-5, 2e-3
        torch.testing.assert_allclose(out, ref, atol=atol, rtol=rtol)
//  [python]
def test_gelu(self, dtype, strongly_typed):
        if dtype == 'bfloat16':
            self.skip_bf16_before_ampere()

        torch_dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype)
        x_shape = (12, 12, 96, 96)
        x_data = torch.rand(x_shape, dtype=torch_dtype)

        # construct trt network
        builder = tensorrt_llm.Builder()
        builder.strongly_typed = strongly_typed
        net = builder.create_network()
        with tensorrt_llm.net_guard(net):
            network = tensorrt_llm.default_trtnet()
            x = Tensor(name='x',
                       shape=x_shape,
                       dtype=tensorrt_llm.str_dtype_to_trt(dtype))
            output = tensorrt_llm.functional.gelu(x).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={'x': x_data})
        out = outputs['output'].to(torch_dtype)

        # Reference
        ref = self.gelu(x_data, dtype)

        if dtype == 'bfloat16':
            atol, rtol = 1e-5, 2e-2
        else:
            atol, rtol = 1e-5, 2e-3
        torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
//  [python]
def test_torchvision_numpy_transforms_equivalency(self):
        # as we had to reimplement the torchvision transforms using transformers utils we must check
        # they both do the same

        image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False)
        image_processor = self.image_processing_class(**self.image_processor_dict)

        print(image_inputs)

        def convert_to_rgb(image):
            # `image.convert("RGB")` would only work for .jpg images, as it creates a wrong background
            # for transparent images. The call to `alpha_composite` handles this case
            if image.mode == "RGB":
                return image

            image_rgba = image.convert("RGBA")
            background = Image.new("RGBA", image_rgba.size, (255, 255, 255))
            alpha_composite = Image.alpha_composite(background, image_rgba)
            alpha_composite = alpha_composite.convert("RGB")
            return alpha_composite

        image_size = image_processor.image_size
        image_mean = image_processor.image_mean
        image_std = image_processor.image_std

        transform = transforms.Compose(
            [
                convert_to_rgb,
                transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC),
                transforms.ToTensor(),
                transforms.Normalize(mean=image_mean, std=image_std),
            ]
        )

        pixel_values_transform_implied = image_processor(image_inputs, transform=None)
        pixel_values_transform_supplied = image_processor(image_inputs, transform=transform)

        torch.testing.assert_allclose(pixel_values_transform_implied, pixel_values_transform_supplied, rtol=0.0, atol=0.0)
//  [python]
def test_torchvision_numpy_transforms_equivalency(self):
        # as we had to reimplement the torchvision transforms using transformers utils we must check
        # they both do the same

        image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False)
        image_processor = self.image_processing_class(**self.image_processor_dict)

        print(image_inputs)

        def convert_to_rgb(image):
            # `image.convert("RGB")` would only work for .jpg images, as it creates a wrong background
            # for transparent images. The call to `alpha_composite` handles this case
            if image.mode == "RGB":
                return image

            image_rgba = image.convert("RGBA")
            background = Image.new("RGBA", image_rgba.size, (255, 255, 255))
            alpha_composite = Image.alpha_composite(background, image_rgba)
            alpha_composite = alpha_composite.convert("RGB")
            return alpha_composite

        image_size = image_processor.image_size
        image_mean = image_processor.image_mean
        image_std = image_processor.image_std

        transform = transforms.Compose(
            [
                convert_to_rgb,
                transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC),
                transforms.ToTensor(),
                transforms.Normalize(mean=image_mean, std=image_std),
            ]
        )

        pixel_values_transform_implied = image_processor(image_inputs, transform=None)
        pixel_values_transform_supplied = image_processor(image_inputs, transform=transform)

        torch.testing.assert_close(pixel_values_transform_implied, pixel_values_transform_supplied, rtol=0.0, atol=0.0)
//  [python]
def test_isin_1dim(key, dim):
    td = TensorDict(
        {
            "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [7, 8, 9]]),
            "tensor2": torch.tensor([[10, 20], [30, 40], [40, 50], [50, 60]]),
        },
        batch_size=[4],
    )
    td_ref = TensorDict(
        {
            "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [10, 11, 12]]),
            "tensor2": torch.tensor([[10, 20], [30, 40], [50, 60]]),
        },
        batch_size=[3],
    )

    if key == "tensor3":
        with pytest.raises(
            KeyError, match=f"Key '{key}' not found in input or not a tensor."
        ):
            isin(td, td_ref, key, dim)
    elif dim in (1, -2):
        with pytest.raises(
            ValueError,
            match=f"The specified dimension '{dim}' is invalid for an input TensorDict with batch size .*.",
        ):
            isin(td, td_ref, key, dim)
    else:
        in_ref = isin(td, td_ref, key, dim)
        expected_in_ref = torch.tensor([True, True, True, False])
        torch.testing.assert_allclose(in_ref, expected_in_ref)

        with pytest.raises(
            ValueError,
            match="The number of dimensions in the batch size of the input and reference must be the same.",
        ):
            td.batch_size = []
            isin(td, td_ref, key, dim)
//  [python]
def test_isin_1dim(key, dim):
    td = TensorDict(
        {
            "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [7, 8, 9]]),
            "tensor2": torch.tensor([[10, 20], [30, 40], [40, 50], [50, 60]]),
        },
        batch_size=[4],
    )
    td_ref = TensorDict(
        {
            "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [10, 11, 12]]),
            "tensor2": torch.tensor([[10, 20], [30, 40], [50, 60]]),
        },
        batch_size=[3],
    )

    if key == "tensor3":
        with pytest.raises(
            KeyError, match=f"Key '{key}' not found in input or not a tensor."
        ):
            isin(td, td_ref, key, dim)
    elif dim in (1, -2):
        with pytest.raises(
            ValueError,
            match=f"The specified dimension '{dim}' is invalid for an input TensorDict with batch size .*.",
        ):
            isin(td, td_ref, key, dim)
    else:
        in_ref = isin(td, td_ref, key, dim)
        expected_in_ref = torch.tensor([True, True, True, False])
        torch.testing.assert_close(in_ref, expected_in_ref)

        with pytest.raises(
            ValueError,
            match="The number of dimensions in the batch size of the input and reference must be the same.",
        ):
            td.batch_size = []
            isin(td, td_ref, key, dim)
//  [python]
def test_fb_fc_packed(self):
        X = np.random.rand(16, 16).astype(np.float32) - 0.5
        W = np.random.rand(16, 16).astype(np.float32) - 0.5
        b = np.random.rand(16).astype(np.float32) - 0.5

        def fc_op(X, W, b):
            return np.dot(X, W.T) + b

        x_tensor = torch.tensor(X)
        w_tensor = torch.tensor(W)
        b_tensor = torch.tensor(b)
        packed_w_tensor = torch.fbgemm_pack_gemm_matrix_fp16(w_tensor)
        actual_output = torch.fbgemm_linear_fp16_weight(x_tensor, packed_w_tensor, b_tensor)
        expected_output = fc_op(X, W, b)
        torch.testing.assert_allclose(torch.from_numpy(expected_output), actual_output.cpu(), atol=1e-3, rtol=1e-3)
//  [python]
def test_fb_fc_packed(self):
        X = np.random.rand(16, 16).astype(np.float32) - 0.5
        W = np.random.rand(16, 16).astype(np.float32) - 0.5
        b = np.random.rand(16).astype(np.float32) - 0.5

        def fc_op(X, W, b):
            return np.dot(X, W.T) + b

        x_tensor = torch.tensor(X)
        w_tensor = torch.tensor(W)
        b_tensor = torch.tensor(b)
        packed_w_tensor = torch.fbgemm_pack_gemm_matrix_fp16(w_tensor)
        actual_output = torch.fbgemm_linear_fp16_weight(x_tensor, packed_w_tensor, b_tensor)
        expected_output = fc_op(X, W, b)
        torch.testing.assert_close(torch.from_numpy(expected_output), actual_output.cpu(), atol=1e-3, rtol=1e-3)
//  [python]
def _test_jagged_dense_dense_elementwise_add_jagged_output(
        self,
        num_jagged_dim: int,
        outer_dense_size: int,
        inner_dense_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        x_values, x_offsets, max_lengths = generate_jagged_tensor(
            num_jagged_dim, outer_dense_size, inner_dense_size, dtype, device
        )

        x_padded = to_padded_dense(x_values, x_offsets, max_lengths)
        # create a jagged tensor and then densify
        y_0 = to_padded_dense(
            torch.rand(
                (
                    max(outer_dense_size * np.prod(max_lengths), x_values.size(0)),
                    inner_dense_size,
                ),
                dtype=dtype,
                device=device,
            ),
            x_offsets,
            max_lengths,
        )
        y_1 = to_padded_dense(
            torch.rand(
                (
                    max(outer_dense_size * np.prod(max_lengths), x_values.size(0)),
                    inner_dense_size,
                ),
                dtype=dtype,
                device=device,
            ),
            x_offsets,
            max_lengths,
        )
        output_ref = x_padded + y_0 + y_1
        (
            output,
            output_offsets,
        ) = torch.ops.fbgemm.jagged_dense_dense_elementwise_add_jagged_output(
            x_values, x_offsets, y_0, y_1
        )
        output = to_padded_dense(output, output_offsets, max_lengths)

        torch.testing.assert_allclose(output, output_ref)

        # pyre-fixme[2]: Parameter must be annotated.
        def add_jagged_output_func(*args) -> torch.Tensor:
            return torch.ops.fbgemm.jagged_dense_dense_elementwise_add_jagged_output(
                *args
            )[0]

        f = add_jagged_output_func

        gradcheck(
            f,
            (
                x_values.float().requires_grad_(True),
                x_offsets,
                y_0.float().requires_grad_(True),
                y_1.float().requires_grad_(True),
            ),
            eps=1e-2,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def _test_jagged_dense_dense_elementwise_add_jagged_output(
        self,
        num_jagged_dim: int,
        outer_dense_size: int,
        inner_dense_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        x_values, x_offsets, max_lengths = generate_jagged_tensor(
            num_jagged_dim, outer_dense_size, inner_dense_size, dtype, device
        )

        x_padded = to_padded_dense(x_values, x_offsets, max_lengths)
        # create a jagged tensor and then densify
        y_0 = to_padded_dense(
            torch.rand(
                (
                    max(outer_dense_size * np.prod(max_lengths), x_values.size(0)),
                    inner_dense_size,
                ),
                dtype=dtype,
                device=device,
            ),
            x_offsets,
            max_lengths,
        )
        y_1 = to_padded_dense(
            torch.rand(
                (
                    max(outer_dense_size * np.prod(max_lengths), x_values.size(0)),
                    inner_dense_size,
                ),
                dtype=dtype,
                device=device,
            ),
            x_offsets,
            max_lengths,
        )
        output_ref = x_padded + y_0 + y_1
        (
            output,
            output_offsets,
        ) = torch.ops.fbgemm.jagged_dense_dense_elementwise_add_jagged_output(
            x_values, x_offsets, y_0, y_1
        )
        output = to_padded_dense(output, output_offsets, max_lengths)

        torch.testing.assert_close(output, output_ref)

        # pyre-fixme[2]: Parameter must be annotated.
        def add_jagged_output_func(*args) -> torch.Tensor:
            return torch.ops.fbgemm.jagged_dense_dense_elementwise_add_jagged_output(
                *args
            )[0]

        f = add_jagged_output_func

        gradcheck(
            f,
            (
                x_values.float().requires_grad_(True),
                x_offsets,
                y_0.float().requires_grad_(True),
                y_1.float().requires_grad_(True),
            ),
            eps=1e-2,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def test_psnr_skimage_equivelant(self) -> None:
        input, target = self._get_random_data_peak_signal_to_noise_ratio(
            BATCH_SIZE, IMG_CHANNELS, IMG_HEIGHT, IMG_WIDTH
        )

        input_np = input.numpy().ravel()
        target_np = target.numpy().ravel()
        skimage_result = torch.tensor(
            skimage_psnr(target_np, input_np), dtype=torch.float32
        )

        torch.testing.assert_allclose(
            peak_signal_noise_ratio(input, target),
            skimage_result,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def test_psnr_skimage_equivelant(self) -> None:
        input, target = self._get_random_data_peak_signal_to_noise_ratio(
            BATCH_SIZE, IMG_CHANNELS, IMG_HEIGHT, IMG_WIDTH
        )

        input_np = input.numpy().ravel()
        target_np = target.numpy().ravel()
        skimage_result = torch.tensor(
            skimage_psnr(target_np, input_np), dtype=torch.float32
        )

        torch.testing.assert_close(
            peak_signal_noise_ratio(input, target),
            skimage_result,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def test_index_select_dim0(
        self,
        N: int,
        shape: List[int],
        dtype: torch.dtype,
        use_cpu: bool,
        consecutive_indices: bool,
        skip_indices_sorting_fwd: bool,
        use_inference_mode: bool,
    ) -> None:
        device = torch.device("cpu" if use_cpu else "cuda")
        U = random.randint(0, N + 1)

        kwargs = {}
        if consecutive_indices:
            start = np.random.randint(0, U)
            length = np.random.randint(1, U - start + 1)
            indices = list(range(start, start + length))
            np_arr = np.array(indices)
            for _ in range(N - U):
                indices.append(np.random.randint(start, start + length))
                np_arr = np.array(indices)
                np.random.shuffle(np_arr)
            indices = torch.from_numpy(np_arr).to(torch.int).to(device)
            kwargs["consecutive_range_start"] = start
            kwargs["consecutive_range_length"] = length
        else:
            indices = torch.randint(U, (N,), device=device)

        kwargs["skip_indices_sorting_fwd"] = skip_indices_sorting_fwd

        input = torch.rand((U,) + tuple(shape), dtype=dtype, device=device)

        with torch.inference_mode() if use_inference_mode else contextlib.nullcontext():
            output_ref = torch.ops.fbgemm.index_select_dim0(input, indices, **kwargs)
            output = torch.index_select(input, 0, indices)

            torch.testing.assert_allclose(output, output_ref)

        if not use_inference_mode:
            gradcheck_args = [
                input.clone().detach().float().requires_grad_(True),
                indices,
            ]
            for k in kwargs:
                gradcheck_args.append(kwargs[k])

            torch.autograd.gradcheck(
                torch.ops.fbgemm.index_select_dim0,
                gradcheck_args,
                eps=1e-2,
                atol=1e-3,
                rtol=1e-3,
            )
//  [python]
def test_index_select_dim0(
        self,
        N: int,
        shape: List[int],
        dtype: torch.dtype,
        use_cpu: bool,
        consecutive_indices: bool,
        skip_indices_sorting_fwd: bool,
        use_inference_mode: bool,
    ) -> None:
        device = torch.device("cpu" if use_cpu else "cuda")
        U = random.randint(0, N + 1)

        kwargs = {}
        if consecutive_indices:
            start = np.random.randint(0, U)
            length = np.random.randint(1, U - start + 1)
            indices = list(range(start, start + length))
            np_arr = np.array(indices)
            for _ in range(N - U):
                indices.append(np.random.randint(start, start + length))
                np_arr = np.array(indices)
                np.random.shuffle(np_arr)
            indices = torch.from_numpy(np_arr).to(torch.int).to(device)
            kwargs["consecutive_range_start"] = start
            kwargs["consecutive_range_length"] = length
        else:
            indices = torch.randint(U, (N,), device=device)

        kwargs["skip_indices_sorting_fwd"] = skip_indices_sorting_fwd

        input = torch.rand((U,) + tuple(shape), dtype=dtype, device=device)

        with torch.inference_mode() if use_inference_mode else contextlib.nullcontext():
            output_ref = torch.ops.fbgemm.index_select_dim0(input, indices, **kwargs)
            output = torch.index_select(input, 0, indices)

            torch.testing.assert_close(output, output_ref)

        if not use_inference_mode:
            gradcheck_args = [
                input.clone().detach().float().requires_grad_(True),
                indices,
            ]
            for k in kwargs:
                gradcheck_args.append(kwargs[k])

            torch.autograd.gradcheck(
                torch.ops.fbgemm.index_select_dim0,
                gradcheck_args,
                eps=1e-2,
                atol=1e-3,
                rtol=1e-3,
            )
//  [python]
def test_tensor_index_slice(self):
        def fn(a):
            x = torch.tensor([1, 2], device=self.device)
            y = torch.tensor([2, 3], device=self.device)
            xx = torch.tensor([1, 2], device=self.device).view(1, 2)
            yy = torch.tensor([1, 2, 3], device=self.device).view(3, 1)
            return [
                a[x, y],
                a[:, x, y],
                a[:, x, y, :],
                a[x, :, y],
                a[:, x, :, y, :],
                a[xx, yy],
                a[:, xx, yy],
                a[xx, :, yy],
                a[xx, yy, :],
                a[:, xx, :, yy],
            ]

        a = torch.arange(3 * 4 * 5 * 6 * 7, device=self.device).view(3, 4, 5, 6, 7)
        refs = fn(a)
        tests = torch.compile(fn)(a)
        for ref, test in zip(refs, tests):
            self.assertEqual(ref, test)
//  [python]
def test_tensor_index_slice(self):
        def fn(a):
            x = torch.tensor([1, 2], device=self.device)
            y = torch.tensor([2, 3], device=self.device)
            xx = torch.tensor([1, 2], device=self.device).view(1, 2)
            yy = torch.tensor([1, 2, 3], device=self.device).view(3, 1)
            return [
                a[x, y],
                a[:, x, y],
                a[:, x, y, :],
                a[x, :, y],
                a[:, x, :, y, :],
                a[xx, yy],
                a[:, xx, yy],
                a[xx, :, yy],
                a[xx, yy, :],
                a[:, xx, :, yy],
            ]

        a = torch.arange(3 * 4 * 5 * 6 * 7, device=self.device).view(3, 4, 5, 6, 7)
        refs = fn(a)
        tests = torch.compile(fn)(a)
        for ref, test in zip(refs, tests):
            torch.testing.assert_close(ref, test)
//  [python]
def _test_remove_elementwise_no_ops_impl(
        self,
        elementwise_op_getter: Callable[[Tensor], Tensor],
        expected_op: Callable[[Tensor], Tensor],
    ):
        dtype = "float16"
        batch_dim = shape_utils.gen_int_var_min_max(self.BATCH_SIZES, "batch_0")
        reduce_dim = 0
        X0 = gen_input_tensor([batch_dim, IntImm(self.M)], name="x0", dtype=dtype)
        elementwise_op_0 = elementwise_op_getter(X0)
        Y = ops.reduce_mean(reduce_dim)(elementwise_op_0)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        module = compile_model(
            Y,
            detect_target(),
            "./tmp",
            "test_remove_elementwise_no_ops",
        )

        sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
        self.assertEqual(len(sorted_ops), 1)

        for batch in self.BATCH_SIZES:
            x0_pt = get_random_torch_tensor([batch, self.M], dtype)
            add_0_pt = expected_op(x0_pt)
            y_pt = torch.mean(add_0_pt, dim=reduce_dim)
            y = get_torch_empty_tensor(y_pt.size(), dtype)
            inputs = {"x0": x0_pt}
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(y_pt, y, atol=0.01, rtol=0.01)
//  [python]
def _test_remove_elementwise_no_ops_impl(
        self,
        elementwise_op_getter: Callable[[Tensor], Tensor],
        expected_op: Callable[[Tensor], Tensor],
    ):
        dtype = "float16"
        batch_dim = shape_utils.gen_int_var_min_max(self.BATCH_SIZES, "batch_0")
        reduce_dim = 0
        X0 = gen_input_tensor([batch_dim, IntImm(self.M)], name="x0", dtype=dtype)
        elementwise_op_0 = elementwise_op_getter(X0)
        Y = ops.reduce_mean(reduce_dim)(elementwise_op_0)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        module = compile_model(
            Y,
            detect_target(),
            "./tmp",
            "test_remove_elementwise_no_ops",
        )

        sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
        self.assertEqual(len(sorted_ops), 1)

        for batch in self.BATCH_SIZES:
            x0_pt = get_random_torch_tensor([batch, self.M], dtype)
            add_0_pt = expected_op(x0_pt)
            y_pt = torch.mean(add_0_pt, dim=reduce_dim)
            y = get_torch_empty_tensor(y_pt.size(), dtype)
            inputs = {"x0": x0_pt}
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y, atol=0.01, rtol=0.01)
//  [python]
def test_decompose():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 20)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = self.linear(x)
            x = self.relu(x)
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    sch["linear"].decompose()
    sch.trace(flatten=True)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((32, 10), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_allclose(out, out_ref)
//  [python]
def test_decompose():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 20)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = self.linear(x)
            x = self.relu(x)
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    sch["linear"].decompose()
    sch.trace(flatten=True)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((32, 10), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_vertical_fusion():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = nn.Conv2d(3, 32, 3)

        def forward(self, x):
            x = self.conv(x)
            x = F.relu(x)
            x = x + 1
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    def pattern(x: torch.Tensor):
        x = call_module("conv", x)
        x = F.relu(x)
        x = x + 1
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 3
    sch.fuse(subgraph, compiler="TorchScript", name="FusedReLU")

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 3, 32, 32), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_allclose(out, out_ref)
//  [python]
def test_vertical_fusion():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = nn.Conv2d(3, 32, 3)

        def forward(self, x):
            x = self.conv(x)
            x = F.relu(x)
            x = x + 1
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    def pattern(x: torch.Tensor):
        x = call_module("conv", x)
        x = F.relu(x)
        x = x + 1
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 3
    sch.fuse(subgraph, compiler="TorchScript", name="FusedReLU")

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 3, 32, 32), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_bias_gelu():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(1024, 1024)
            self.gelu = nn.GELU()

        def forward(self, x):
            x = self.linear(x)
            x = self.gelu(x)
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    sch["linear"].decompose()
    sch.trace(flatten=True)
    print(sch.mod.graph)

    def pattern(x, bias):
        x = F.gelu(bias + x)
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 2
    assert subgraph[0][0][1].target == operator.add
    assert subgraph[0][1][1].target == "gelu"
    sch.fuse(subgraph, compiler="TorchScript", name="BiasGeLU")
    assert isinstance(sch["BiasGeLU_0"].mod, torch.jit.ScriptModule)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    t.assert_close(out, out_ref)