jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testBatchNorm(self):
        device = get_test_device()
        x = torch.randn(16, 3, 224, 224, device=device)
        weight = torch.randn(3, device=device)
        bias = torch.randn(3, device=device)

        for i in range(10):
            # BatchNorm2d does extra checks on dimensions which SymInts don't support yet
            # so we call `torch.ops.aten.native_batch_norm` to bypass the checks.
            z, _, _ = torch.ops.aten.native_batch_norm(
                x, weight, bias, None, None, True, 0.1, 1e-5
            )
            z_legit, _, _ = torch.ops.aten._native_batch_norm_legit(
                x, weight, bias, True, 0.1, 1e-5
            )

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        weight_lazy = weight.detach().clone().to(device=device)
        bias_lazy = bias.detach().clone().to(device=device)
        for i in range(10):
            z_lazy, _, _ = torch.ops.aten.native_batch_norm(
                x_lazy, weight_lazy, bias_lazy, None, None, True, 0.1, 1e-5
            )
            z_legit_lazy, _, _ = torch.ops.aten._native_batch_norm_legit(
                x_lazy, weight_lazy, bias_lazy, True, 0.1, 1e-5
            )
            torch._lazy.mark_step()

        torch.testing.assert_allclose(z.cpu(), z_lazy.cpu())
        torch.testing.assert_allclose(z_legit.cpu(), z_legit_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::NativeBatchNorm") >= 7
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
def testBatchNorm(self):
        device = get_test_device()
        x = torch.randn(16, 3, 224, 224, device=device)
        weight = torch.randn(3, device=device)
        bias = torch.randn(3, device=device)

        for i in range(10):
            # BatchNorm2d does extra checks on dimensions which SymInts don't support yet
            # so we call `torch.ops.aten.native_batch_norm` to bypass the checks.
            z, _, _ = torch.ops.aten.native_batch_norm(
                x, weight, bias, None, None, True, 0.1, 1e-5
            )
            z_legit, _, _ = torch.ops.aten._native_batch_norm_legit(
                x, weight, bias, True, 0.1, 1e-5
            )

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        weight_lazy = weight.detach().clone().to(device=device)
        bias_lazy = bias.detach().clone().to(device=device)
        for i in range(10):
            z_lazy, _, _ = torch.ops.aten.native_batch_norm(
                x_lazy, weight_lazy, bias_lazy, None, None, True, 0.1, 1e-5
            )
            z_legit_lazy, _, _ = torch.ops.aten._native_batch_norm_legit(
                x_lazy, weight_lazy, bias_lazy, True, 0.1, 1e-5
            )
            torch._lazy.mark_step()

        torch.testing.assert_close(z.cpu(), z_lazy.cpu())
        torch.testing.assert_close(z_legit.cpu(), z_legit_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::NativeBatchNorm") >= 7
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
def test_rope():
    bs, seq_len, n_head, n_embed = 1, 6, 2, 8
    head_size = n_embed // n_head
    x = torch.randint(0, 10000, size=(bs, n_head, seq_len, head_size)).float()
    position_ids = torch.arange(seq_len).unsqueeze(0)

    theirs = GPTNeoXRotaryEmbedding(head_size, seq_len)
    ours_cos_cached, ours_sin_cached = build_rope_cache(seq_len, head_size, device=x.device)
    # their rope cache has 2 added dimensions and the cos/sin is duplicated
    torch.testing.assert_allclose(ours_cos_cached, theirs.cos_cached.squeeze())
    torch.testing.assert_allclose(ours_sin_cached, theirs.sin_cached.squeeze())

    ours_x_rope = apply_rope(x, ours_cos_cached, ours_sin_cached)
    theirs_x_rope, _ = apply_rotary_pos_emb(x, x, theirs.cos_cached, theirs.sin_cached, position_ids)
    torch.testing.assert_allclose(ours_x_rope, theirs_x_rope)
//  [python]
def test_rope():
    bs, seq_len, n_head, n_embed = 1, 6, 2, 8
    head_size = n_embed // n_head
    x = torch.randint(0, 10000, size=(bs, n_head, seq_len, head_size)).float()
    position_ids = torch.arange(seq_len).unsqueeze(0)

    theirs = GPTNeoXRotaryEmbedding(head_size, seq_len)
    ours_cos_cached, ours_sin_cached = build_rope_cache(seq_len, head_size, device=x.device)
    # their rope cache has 2 added dimensions and the cos/sin is duplicated
    torch.testing.assert_close(ours_cos_cached, theirs.cos_cached.squeeze())
    torch.testing.assert_close(ours_sin_cached, theirs.sin_cached.squeeze())

    ours_x_rope = apply_rope(x, ours_cos_cached, ours_sin_cached)
    theirs_x_rope, _ = apply_rotary_pos_emb(x, x, theirs.cos_cached, theirs.sin_cached, position_ids)
    torch.testing.assert_close(ours_x_rope, theirs_x_rope)
//  [python]
def _test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=1e-8, **fn_kwargs):
    transformed_batch = fn(batch_tensors, **fn_kwargs)
    for i in range(len(batch_tensors)):
        img_tensor = batch_tensors[i, ...]
        transformed_img = fn(img_tensor, **fn_kwargs)
        torch.testing.assert_allclose(transformed_img, transformed_batch[i, ...], rtol=0, atol=1e-6)

    if scripted_fn_atol >= 0:
        scripted_fn = torch.jit.script(fn)
        # scriptable function test
        s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs)
        torch.testing.assert_allclose(transformed_batch, s_transformed_batch, rtol=1e-5, atol=scripted_fn_atol)
//  [python]
def _test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=1e-8, **fn_kwargs):
    transformed_batch = fn(batch_tensors, **fn_kwargs)
    for i in range(len(batch_tensors)):
        img_tensor = batch_tensors[i, ...]
        transformed_img = fn(img_tensor, **fn_kwargs)
        torch.testing.assert_close(transformed_img, transformed_batch[i, ...], rtol=0, atol=1e-6)

    if scripted_fn_atol >= 0:
        scripted_fn = torch.jit.script(fn)
        # scriptable function test
        s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs)
        torch.testing.assert_close(transformed_batch, s_transformed_batch, rtol=1e-5, atol=scripted_fn_atol)
//  [python]
def test_log_prob():
    gibbs_dag_prior = GibbsDAGPrior(num_nodes=123, sparsity_lambda=torch.tensor(1))

    A = torch.Tensor([[0, 0, 1], [0, 0, 1], [0, 0, 0]])

    with pytest.raises(AssertionError):
        gibbs_dag_prior.log_prob(A)

    gibbs_dag_prior = GibbsDAGPrior(num_nodes=2, sparsity_lambda=torch.tensor(1))

    A = torch.Tensor([[1, 1], [0, 1]])

    torch.testing.assert_allclose(
        gibbs_dag_prior.log_prob(A),
        torch.tensor(-3.0),
    )

    A = torch.Tensor([[0, 1], [0, 0]])

    torch.testing.assert_allclose(gibbs_dag_prior.log_prob(A), torch.tensor(-1.0))
//  [python]
def test_log_prob():
    gibbs_dag_prior = GibbsDAGPrior(num_nodes=123, sparsity_lambda=torch.tensor(1))

    A = torch.Tensor([[0, 0, 1], [0, 0, 1], [0, 0, 0]])

    with pytest.raises(AssertionError):
        gibbs_dag_prior.log_prob(A)

    gibbs_dag_prior = GibbsDAGPrior(num_nodes=2, sparsity_lambda=torch.tensor(1))

    A = torch.Tensor([[1, 1], [0, 1]])

    torch.testing.assert_close(
        gibbs_dag_prior.log_prob(A),
        torch.tensor(-3.0),
    )

    A = torch.Tensor([[0, 1], [0, 0]])

    torch.testing.assert_close(gibbs_dag_prior.log_prob(A), torch.tensor(-1.0))
//  [python]
def align_tensors(tensor_alignment_data_iter: Iterable[TensorAlignmentData]):
    """
    Checks the alignment between tensors specified by
    ``tensor_alignment_data_iter``. Each element in the iterable specifies a
    single tensor quantity to align between FlexFlow and PyTorch.
    """
    for tensor_alignment_data in tensor_alignment_data_iter:
        ff_filepath = tensor_alignment_data.ff_filepath
        torch_filepath = tensor_alignment_data.torch_filepath
        assert os.path.exists(ff_filepath), \
            f"Missing FlexFlow tensor at {ff_filepath}"
        assert os.path.exists(torch_filepath), \
            f"Missing PyTorch tensor at {torch_filepath}"
        ff_tensor = torch.load(ff_filepath).cpu()
        torch_tensor = torch.load(torch_filepath).cpu()
        print(f"Checking {tensor_alignment_data.tensor_name} alignment...")
        torch.testing.assert_allclose(ff_tensor, torch_tensor, rtol=1e-2, atol=1e-4)
//  [python]
def align_tensors(tensor_alignment_data_iter: Iterable[TensorAlignmentData]):
    """
    Checks the alignment between tensors specified by
    ``tensor_alignment_data_iter``. Each element in the iterable specifies a
    single tensor quantity to align between FlexFlow and PyTorch.
    """
    for tensor_alignment_data in tensor_alignment_data_iter:
        ff_filepath = tensor_alignment_data.ff_filepath
        torch_filepath = tensor_alignment_data.torch_filepath
        assert os.path.exists(ff_filepath), \
            f"Missing FlexFlow tensor at {ff_filepath}"
        assert os.path.exists(torch_filepath), \
            f"Missing PyTorch tensor at {torch_filepath}"
        ff_tensor = torch.load(ff_filepath).cpu()
        torch_tensor = torch.load(torch_filepath).cpu()
        print(f"Checking {tensor_alignment_data.tensor_name} alignment...")
        torch.testing.assert_close(ff_tensor, torch_tensor, rtol=1e-2, atol=1e-4)
//  [python]
def test_get_expert_graph_term():

    mask = torch.Tensor([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    dag = torch.Tensor([[0, 0, 1], [0, 0, 0], [0, 0, 0]])
    confidence = 0.8
    scale = 10

    expert_graph_container = ExpertGraphContainer(dag, mask, confidence, scale)

    gibbs_dag_prior = GibbsDAGPrior(
        num_nodes=3,
        sparsity_lambda=torch.tensor(1),
        expert_graph_container=expert_graph_container,
    )

    A = torch.Tensor([[0, 0, 1], [0, 1, 1], [1, 0, 1]])

    assert gibbs_dag_prior.get_expert_graph_term(A) == torch.tensor(0)

    mask = torch.Tensor([[0, 0, 1], [0, 0, 0], [0, 0, 0]])

    expert_graph_container = ExpertGraphContainer(dag, mask, confidence, scale)
    gibbs_dag_prior = GibbsDAGPrior(
        num_nodes=3,
        sparsity_lambda=torch.tensor(1),
        expert_graph_container=expert_graph_container,
    )

    t.assert_allclose(gibbs_dag_prior.get_expert_graph_term(A), torch.tensor(0.2))
//  [python]
def test_get_expert_graph_term():

    mask = torch.Tensor([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    dag = torch.Tensor([[0, 0, 1], [0, 0, 0], [0, 0, 0]])
    confidence = 0.8
    scale = 10

    expert_graph_container = ExpertGraphContainer(dag, mask, confidence, scale)

    gibbs_dag_prior = GibbsDAGPrior(
        num_nodes=3,
        sparsity_lambda=torch.tensor(1),
        expert_graph_container=expert_graph_container,
    )

    A = torch.Tensor([[0, 0, 1], [0, 1, 1], [1, 0, 1]])

    assert gibbs_dag_prior.get_expert_graph_term(A) == torch.tensor(0)

    mask = torch.Tensor([[0, 0, 1], [0, 0, 0], [0, 0, 0]])

    expert_graph_container = ExpertGraphContainer(dag, mask, confidence, scale)
    gibbs_dag_prior = GibbsDAGPrior(
        num_nodes=3,
        sparsity_lambda=torch.tensor(1),
        expert_graph_container=expert_graph_container,
    )

    torch.testing.assert_close(gibbs_dag_prior.get_expert_graph_term(A), torch.tensor(0.2))
//  [python]
def test_padded_collate(self):
        """
        Tests that shorter input, label sequences are padded to the max seq len.
        """
        padding_idx = -8
        ignore_idx = -9
        token_pairs = [
            ([1, 2, 3], [4, 5, 6]),
            (
                [
                    7,
                ],
                [
                    10,
                ],
            ),
        ]
        inputs, labels = padded_collate(
            batch=token_pairs,
            padding_idx=padding_idx,
            ignore_idx=ignore_idx,
        )
        padded_input = inputs[1]
        padded_label = labels[1]
        t.assert_allclose(
            padded_input, torch.tensor([7, padding_idx, padding_idx])
        )
        t.assert_allclose(
            padded_label, torch.tensor([10, ignore_idx, ignore_idx])
        )
//  [python]
def test_padded_collate(self):
        """
        Tests that shorter input, label sequences are padded to the max seq len.
        """
        padding_idx = -8
        ignore_idx = -9
        token_pairs = [
            ([1, 2, 3], [4, 5, 6]),
            (
                [
                    7,
                ],
                [
                    10,
                ],
            ),
        ]
        inputs, labels = padded_collate(
            batch=token_pairs,
            padding_idx=padding_idx,
            ignore_idx=ignore_idx,
        )
        padded_input = inputs[1]
        padded_label = labels[1]
        torch.testing.assert_close(
            padded_input, torch.tensor([7, padding_idx, padding_idx])
        )
        torch.testing.assert_close(
            padded_label, torch.tensor([10, ignore_idx, ignore_idx])
        )
//  [python]
def test_dataloader_from_tarred_nemo_manifest_concat(nemo_tarred_manifest_path: tuple[str, str]):
    json_mft, tar_mft = nemo_tarred_manifest_path
    config = OmegaConf.create(
        {
            "manifest_filepath": json_mft,
            "tarred_audio_filepaths": tar_mft,
            "sample_rate": 16000,
            "shuffle": True,
            "use_lhotse": True,
            "num_workers": 0,
            # lhotse specific
            "concatenate_samples": True,
            "concatenate_duration_factor": 3.0,
            "batch_duration": 4.0,
            "quadratic_duration": 15.0,  # seconds
            "use_bucketing": False,
            "drop_last": False,
            "shuffle_buffer_size": 10,
            "seed": 0,
            "shard_seed": 0,
        }
    )

    dl = get_lhotse_dataloader_from_config(
        config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
    )

    batches = [batch for batch in islice(dl, 4)]

    assert len(batches) == 4

    # the first element has been concatenated: 2x16000 speech (2x1s) + 1600 gap (0.1s)
    expected_audio_lens = torch.tensor([33600, 16000], dtype=torch.int32)

    b = batches[0]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_allclose(b["audio_lens"], expected_audio_lens)

    b = batches[1]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_allclose(b["audio_lens"], expected_audio_lens)

    b = batches[2]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_allclose(b["audio_lens"], expected_audio_lens)

    b = batches[3]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_allclose(b["audio_lens"], expected_audio_lens)
//  [python]
def test_dataloader_from_tarred_nemo_manifest_concat(nemo_tarred_manifest_path: tuple[str, str]):
    json_mft, tar_mft = nemo_tarred_manifest_path
    config = OmegaConf.create(
        {
            "manifest_filepath": json_mft,
            "tarred_audio_filepaths": tar_mft,
            "sample_rate": 16000,
            "shuffle": True,
            "use_lhotse": True,
            "num_workers": 0,
            # lhotse specific
            "concatenate_samples": True,
            "concatenate_duration_factor": 3.0,
            "batch_duration": 4.0,
            "quadratic_duration": 15.0,  # seconds
            "use_bucketing": False,
            "drop_last": False,
            "shuffle_buffer_size": 10,
            "seed": 0,
            "shard_seed": 0,
        }
    )

    dl = get_lhotse_dataloader_from_config(
        config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
    )

    batches = [batch for batch in islice(dl, 4)]

    assert len(batches) == 4

    # the first element has been concatenated: 2x16000 speech (2x1s) + 1600 gap (0.1s)
    expected_audio_lens = torch.tensor([33600, 16000], dtype=torch.int32)

    b = batches[0]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_close(b["audio_lens"], expected_audio_lens)

    b = batches[1]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_close(b["audio_lens"], expected_audio_lens)

    b = batches[2]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_close(b["audio_lens"], expected_audio_lens)

    b = batches[3]
    assert set(b.keys()) == {"audio", "audio_lens", "ids"}
    assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
    torch.testing.assert_close(b["audio_lens"], expected_audio_lens)
//  [python]
def _test_sliced_flat_tensor(parallel_context: ParallelContext):
    a = torch.randn(2, 3, requires_grad=True)
    grad = torch.randn(2, 3)
    a.grad = grad

    start_offset, end_offset = 1, 5
    b = SlicedFlatTensor(a, start_offset=start_offset, end_offset=end_offset)

    t.assert_close(a.grad, grad, atol=0, rtol=0)
    t.assert_close(b.grad, grad.view(-1)[start_offset:end_offset])

    # Deallocate the gradient by setting it to None
    a.grad = None

    assert a.grad is None
    assert b.grad is None

    # Setting gradient to None on the sliced tensor works
    a.grad = grad
    assert a.grad is not None
    assert b.grad is not None
    b.grad = None
    assert b.grad is None
    assert a.grad is None

    with assert_fail_with(NotImplementedError):
        b.grad = torch.randn(1, 5)

    with assert_fail_with(NotImplementedError):
        del b.grad

    c = b[:3]
    # It's important not to contaminate everyone.
    assert not isinstance(c, SlicedFlatTensor)

    parallel_context.destroy()
//  [python]
def _test_sliced_flat_tensor(parallel_context: ParallelContext):
    a = torch.randn(2, 3, requires_grad=True)
    grad = torch.randn(2, 3)
    a.grad = grad

    start_offset, end_offset = 1, 5
    b = SlicedFlatTensor(a, start_offset=start_offset, end_offset=end_offset)

    torch.testing.assert_close(a.grad, grad, atol=0, rtol=0)
    torch.testing.assert_close(b.grad, grad.view(-1)[start_offset:end_offset])

    # Deallocate the gradient by setting it to None
    a.grad = None

    assert a.grad is None
    assert b.grad is None

    # Setting gradient to None on the sliced tensor works
    a.grad = grad
    assert a.grad is not None
    assert b.grad is not None
    b.grad = None
    assert b.grad is None
    assert a.grad is None

    with assert_fail_with(NotImplementedError):
        b.grad = torch.randn(1, 5)

    with assert_fail_with(NotImplementedError):
        del b.grad

    c = b[:3]
    # It's important not to contaminate everyone.
    assert not isinstance(c, SlicedFlatTensor)

    parallel_context.destroy()
//  [python]
def test_fused_rope(
    dtype: torch.dtype,
    seq_length: int,
    hidden_size: int,
    rotary_percent: float,
    margin: int,
    transpose: Union[Tuple, None],
    tensor_format: str,
    loss_func: Callable,
) -> None:
    device = torch.device("cuda:0")
    batch_size, head_num = 2, 64
    t = torch.rand(
        (seq_length - margin, batch_size, head_num, hidden_size),
        dtype=dtype,
        device=device,
    )
    if tensor_format == "bshd":
        t = t.transpose(0, 1).contiguous()
    if transpose:
        t = t.transpose(*transpose).contiguous().transpose(*transpose)
    t.requires_grad = True

    rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent)
    emb = rotary_pos_emb(seq_length)

    # unfused
    output_unfused = apply_rotary_pos_emb(
        t, emb, tensor_format=tensor_format, fused=False
    )
    loss_unfused = loss_func(output_unfused)
    loss_unfused.backward()
    grad_unfused = t.grad.detach().clone()
    t.grad = None

    # fused
    output_fused = apply_rotary_pos_emb(
        t,
        emb,
        tensor_format=tensor_format,
        fused=True,
    )
    loss_fused = loss_func(output_fused)
    loss_fused.backward()
    grad_fused = t.grad.detach().clone()
    t.grad = None

    torch.testing.assert_allclose(output_fused, output_unfused, **get_tol(dtype))
    torch.testing.assert_allclose(grad_fused, grad_unfused, **get_tol(dtype))
    assert output_fused.is_contiguous()
//  [python]
def test_fused_rope(
    dtype: torch.dtype,
    seq_length: int,
    hidden_size: int,
    rotary_percent: float,
    margin: int,
    transpose: Union[Tuple, None],
    tensor_format: str,
    loss_func: Callable,
) -> None:
    device = torch.device("cuda:0")
    batch_size, head_num = 2, 64
    t = torch.rand(
        (seq_length - margin, batch_size, head_num, hidden_size),
        dtype=dtype,
        device=device,
    )
    if tensor_format == "bshd":
        t = t.transpose(0, 1).contiguous()
    if transpose:
        t = t.transpose(*transpose).contiguous().transpose(*transpose)
    t.requires_grad = True

    rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent)
    emb = rotary_pos_emb(seq_length)

    # unfused
    output_unfused = apply_rotary_pos_emb(
        t, emb, tensor_format=tensor_format, fused=False
    )
    loss_unfused = loss_func(output_unfused)
    loss_unfused.backward()
    grad_unfused = t.grad.detach().clone()
    t.grad = None

    # fused
    output_fused = apply_rotary_pos_emb(
        t,
        emb,
        tensor_format=tensor_format,
        fused=True,
    )
    loss_fused = loss_func(output_fused)
    loss_fused.backward()
    grad_fused = t.grad.detach().clone()
    t.grad = None

    torch.testing.assert_close(output_fused, output_unfused, **get_tol(dtype))
    torch.testing.assert_close(grad_fused, grad_unfused, **get_tol(dtype))
    assert output_fused.is_contiguous()
//  [python]
def test_float_to_float(self, input_dtype, output_dtype):
        input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        torch.testing.assert_allclose(output_image_script, output_image, rtol=0.0, atol=1e-6)

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0.0, 1.0

        assert abs(actual_min - desired_min) < 1e-7
        assert abs(actual_max - desired_max) < 1e-7
//  [python]
def test_float_to_float(self, input_dtype, output_dtype):
        input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6)

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0.0, 1.0

        assert abs(actual_min - desired_min) < 1e-7
        assert abs(actual_max - desired_max) < 1e-7
//  [python]
def test_int_to_float(self, input_dtype, output_dtype):
        input_image = torch.tensor((0, torch.iinfo(input_dtype).max), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        torch.testing.assert_allclose(output_image_script, output_image, rtol=0.0, atol=1e-6)

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0.0, 1.0

        assert abs(actual_min - desired_min) < 1e-7
        assert actual_min >= desired_min
        assert abs(actual_max - desired_max) < 1e-7
        assert actual_max <= desired_max
//  [python]
def test_int_to_float(self, input_dtype, output_dtype):
        input_image = torch.tensor((0, torch.iinfo(input_dtype).max), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6)

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0.0, 1.0

        assert abs(actual_min - desired_min) < 1e-7
        assert actual_min >= desired_min
        assert abs(actual_max - desired_max) < 1e-7
        assert actual_max <= desired_max
//  [python]
def test_accimage_to_tensor(self):
        trans = transforms.PILToTensor()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_accimage_to_tensor(self):
        trans = transforms.PILToTensor()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        torch.testing.assert_close(output, expected_output)
//  [python]
def test_float_to_int(self, input_dtype, output_dtype):
        input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        if (input_dtype == torch.float32 and output_dtype in (torch.int32, torch.int64)) or (
            input_dtype == torch.float64 and output_dtype == torch.int64
        ):
            with pytest.raises(RuntimeError):
                transform(input_image)
        else:
            output_image = transform(input_image)
            output_image_script = transform_script(input_image, output_dtype)

            torch.testing.assert_allclose(output_image_script, output_image, rtol=0.0, atol=1e-6)

            actual_min, actual_max = output_image.tolist()
            desired_min, desired_max = 0, torch.iinfo(output_dtype).max

            assert actual_min == desired_min
            assert actual_max == desired_max
//  [python]
def test_float_to_int(self, input_dtype, output_dtype):
        input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        if (input_dtype == torch.float32 and output_dtype in (torch.int32, torch.int64)) or (
            input_dtype == torch.float64 and output_dtype == torch.int64
        ):
            with pytest.raises(RuntimeError):
                transform(input_image)
        else:
            output_image = transform(input_image)
            output_image_script = transform_script(input_image, output_dtype)

            torch.testing.assert_close(output_image_script, output_image, rtol=0.0, atol=1e-6)

            actual_min, actual_max = output_image.tolist()
            desired_min, desired_max = 0, torch.iinfo(output_dtype).max

            assert actual_min == desired_min
            assert actual_max == desired_max
//  [python]
def test_accimage_pil_to_tensor(self):
        trans = transforms.PILToTensor()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_accimage_pil_to_tensor(self):
        trans = transforms.PILToTensor()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        torch.testing.assert_close(output, expected_output)
//  [python]
def test_dtype_int_to_int(self, input_dtype, output_dtype):
        input_max = torch.iinfo(input_dtype).max
        input_image = torch.tensor((0, input_max), dtype=input_dtype)
        output_max = torch.iinfo(output_dtype).max

        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        self.assertEqual(output_image_script, output_image)

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0, output_max

        # see https://github.com/pytorch/vision/pull/2078#issuecomment-641036236 for details
        if input_max >= output_max:
            error_term = 0
        else:
            error_term = 1 - (torch.iinfo(output_dtype).max + 1) // (torch.iinfo(input_dtype).max + 1)

        assert actual_min == desired_min
        assert actual_max == (desired_max + error_term)
//  [python]
def test_dtype_int_to_int(self, input_dtype, output_dtype):
        input_max = torch.iinfo(input_dtype).max
        input_image = torch.tensor((0, input_max), dtype=input_dtype)
        output_max = torch.iinfo(output_dtype).max

        transform = transforms.ConvertImageDtype(output_dtype)
        transform_script = torch.jit.script(F.convert_image_dtype)

        output_image = transform(input_image)
        output_image_script = transform_script(input_image, output_dtype)

        torch.testing.assert_close(
            output_image_script,
            output_image,
            rtol=0.0,
            atol=1e-6,
            msg=f"{output_image_script} vs {output_image}",
        )

        actual_min, actual_max = output_image.tolist()
        desired_min, desired_max = 0, output_max

        # see https://github.com/pytorch/vision/pull/2078#issuecomment-641036236 for details
        if input_max >= output_max:
            error_term = 0
        else:
            error_term = 1 - (torch.iinfo(output_dtype).max + 1) // (torch.iinfo(input_dtype).max + 1)

        assert actual_min == desired_min
        assert actual_max == (desired_max + error_term)
//  [python]
def test_accimage_crop(self):
        trans = transforms.Compose(
            [transforms.CenterCrop(256), transforms.PILToTensor(), transforms.ConvertImageDtype(dtype=torch.float)]
        )

        # Checking if Compose, CenterCrop and ToTensor can be printed as string
        trans.__repr__()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_accimage_crop(self):
        trans = transforms.Compose(
            [transforms.CenterCrop(256), transforms.PILToTensor(), transforms.ConvertImageDtype(dtype=torch.float)]
        )

        # Checking if Compose, CenterCrop and ToTensor can be printed as string
        trans.__repr__()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        torch.testing.assert_close(output, expected_output)
//  [python]
def test_accimage_resize(self):
        trans = transforms.Compose(
            [
                transforms.Resize(256, interpolation=Image.LINEAR),
                transforms.PILToTensor(),
                transforms.ConvertImageDtype(dtype=torch.float),
            ]
        )

        # Checking if Compose, Resize and ToTensor can be printed as string
        trans.__repr__()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        assert np.abs((expected_output - output).mean()) < 1e-3
        assert (expected_output - output).var() < 1e-5
        # note the high absolute tolerance
        torch.testing.assert_allclose(output.numpy(), expected_output.numpy(), rtol=1e-5, atol=5e-2)
//  [python]
def test_accimage_resize(self):
        trans = transforms.Compose(
            [
                transforms.Resize(256, interpolation=Image.LINEAR),
                transforms.PILToTensor(),
                transforms.ConvertImageDtype(dtype=torch.float),
            ]
        )

        # Checking if Compose, Resize and ToTensor can be printed as string
        trans.__repr__()

        expected_output = trans(Image.open(GRACE_HOPPER).convert("RGB"))
        output = trans(accimage.Image(GRACE_HOPPER))

        assert expected_output.size() == output.size()
        assert np.abs((expected_output - output).mean()) < 1e-3
        assert (expected_output - output).var() < 1e-5
        # note the high absolute tolerance
        torch.testing.assert_close(output.numpy(), expected_output.numpy(), rtol=1e-5, atol=5e-2)
//  [python]
def test_autocontrast_equal_minmax():
    img_tensor = torch.tensor([[[10]], [[128]], [[245]]], dtype=torch.uint8).expand(3, 32, 32)
    img_pil = F.to_pil_image(img_tensor)

    img_tensor = F.autocontrast(img_tensor)
    img_pil = F.autocontrast(img_pil)
    torch.testing.assert_allclose(img_tensor, F.pil_to_tensor(img_pil))
//  [python]
def test_autocontrast_equal_minmax():
    img_tensor = torch.tensor([[[10]], [[128]], [[245]]], dtype=torch.uint8).expand(3, 32, 32)
    img_pil = F.to_pil_image(img_tensor)

    img_tensor = F.autocontrast(img_tensor)
    img_pil = F.autocontrast(img_pil)
    torch.testing.assert_close(img_tensor, F.pil_to_tensor(img_pil))
//  [python]
def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        to_tensor = transforms.ToTensor()

        img = transform(img_data)
        assert img.mode == expected_mode
        torch.testing.assert_allclose(expected_output, to_tensor(img).numpy())
//  [python]
def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        to_tensor = transforms.ToTensor()

        img = transform(img_data)
        assert img.mode == expected_mode
        torch.testing.assert_close(expected_output, to_tensor(img).numpy())
//  [python]
def test_pad(self, fill):
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        img = torch.ones(3, height, width, dtype=torch.uint8)
        padding = random.randint(1, 20)
        result = transforms.Compose(
            [
                transforms.ToPILImage(),
                transforms.Pad(padding, fill=fill),
                transforms.PILToTensor(),
            ]
        )(img)
        assert result.size(1) == height + 2 * padding
        assert result.size(2) == width + 2 * padding
        # check that all elements in the padded region correspond
        # to the pad value
        h_padded = result[:, :padding, :]
        w_padded = result[:, :, :padding]
        torch.testing.assert_allclose(h_padded, torch.full_like(h_padded, fill_value=fill), rtol=0.0, atol=0.0)
        torch.testing.assert_allclose(w_padded, torch.full_like(w_padded, fill_value=fill), rtol=0.0, atol=0.0)
        pytest.raises(ValueError, transforms.Pad(padding, fill=(1, 2)), transforms.ToPILImage()(img))
//  [python]
def test_pad(self, fill):
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        img = torch.ones(3, height, width, dtype=torch.uint8)
        padding = random.randint(1, 20)
        result = transforms.Compose(
            [
                transforms.ToPILImage(),
                transforms.Pad(padding, fill=fill),
                transforms.PILToTensor(),
            ]
        )(img)
        assert result.size(1) == height + 2 * padding
        assert result.size(2) == width + 2 * padding
        # check that all elements in the padded region correspond
        # to the pad value
        h_padded = result[:, :padding, :]
        w_padded = result[:, :, :padding]
        torch.testing.assert_close(h_padded, torch.full_like(h_padded, fill_value=fill), rtol=0.0, atol=0.0)
        torch.testing.assert_close(w_padded, torch.full_like(w_padded, fill_value=fill), rtol=0.0, atol=0.0)
        pytest.raises(ValueError, transforms.Pad(padding, fill=(1, 2)), transforms.ToPILImage()(img))
//  [python]
def test_to_tensor(self, channels):
        height, width = 4, 4
        trans = transforms.ToTensor()
        np_rng = np.random.RandomState(0)

        input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        self.assertEqual(output, input_data)

        ndarray = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
        output = trans(ndarray)
        expected_output = ndarray.transpose((2, 0, 1)) / 255.0
        self.assertEqual(output.numpy(), expected_output)

        ndarray = np_rng.rand(height, width, channels).astype(np.float32)
        output = trans(ndarray)
        expected_output = ndarray.transpose((2, 0, 1))
        self.assertEqual(output.numpy(), expected_output)

        # separate test for mode '1' PIL images
        input_data = torch.ByteTensor(1, height, width).bernoulli_()
        img = transforms.ToPILImage()(input_data.mul(255)).convert("1")
        output = trans(img)
        self.assertEqual(input_data, output)
//  [python]
def test_to_tensor(self, channels):
        height, width = 4, 4
        trans = transforms.ToTensor()
        np_rng = np.random.RandomState(0)

        input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        torch.testing.assert_close(output, input_data)

        ndarray = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
        output = trans(ndarray)
        expected_output = ndarray.transpose((2, 0, 1)) / 255.0
        torch.testing.assert_close(output.numpy(), expected_output, check_dtype=False)

        ndarray = np_rng.rand(height, width, channels).astype(np.float32)
        output = trans(ndarray)
        expected_output = ndarray.transpose((2, 0, 1))
        torch.testing.assert_close(output.numpy(), expected_output, check_dtype=False)

        # separate test for mode '1' PIL images
        input_data = torch.ByteTensor(1, height, width).bernoulli_()
        img = transforms.ToPILImage()(input_data.mul(255)).convert("1")
        output = trans(img)
        torch.testing.assert_close(input_data, output, check_dtype=False)
//  [python]
def test_1_channel_float_tensor_to_pil_image(self):
        img_data = torch.Tensor(1, 4, 4).uniform_()
        # 'F' mode for torch.FloatTensor
        img_F_mode = transforms.ToPILImage(mode="F")(img_data)
        assert img_F_mode.mode == "F"
        self.assertEqual(
            np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode)
        )
//  [python]
def test_1_channel_float_tensor_to_pil_image(self):
        img_data = torch.Tensor(1, 4, 4).uniform_()
        # 'F' mode for torch.FloatTensor
        img_F_mode = transforms.ToPILImage(mode="F")(img_data)
        assert img_F_mode.mode == "F"
        torch.testing.assert_close(
            np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode)
        )
//  [python]
def test_pil_to_tensor(self, channels):
        height, width = 4, 4
        trans = transforms.PILToTensor()
        np_rng = np.random.RandomState(0)

        input_data = torch.ByteTensor(channels, height, width).random_(0, 255)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        self.assertEqual(input_data, output)

        input_data = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        expected_output = input_data.transpose((2, 0, 1))
        self.assertEqual(output.numpy(), expected_output)

        input_data = torch.as_tensor(np_rng.rand(channels, height, width).astype(np.float32))
        img = transforms.ToPILImage()(input_data)  # CHW -> HWC and (* 255).byte()
        output = trans(img)  # HWC -> CHW
        expected_output = (input_data * 255).byte()
        self.assertEqual(output, expected_output)

        # separate test for mode '1' PIL images
        input_data = torch.ByteTensor(1, height, width).bernoulli_()
        img = transforms.ToPILImage()(input_data.mul(255)).convert("1")
        output = trans(img).view(torch.uint8).bool().to(torch.uint8)
        self.assertEqual(input_data, output)
//  [python]
def test_pil_to_tensor(self, channels):
        height, width = 4, 4
        trans = transforms.PILToTensor()
        np_rng = np.random.RandomState(0)

        input_data = torch.ByteTensor(channels, height, width).random_(0, 255)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        torch.testing.assert_close(input_data, output)

        input_data = np_rng.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
        img = transforms.ToPILImage()(input_data)
        output = trans(img)
        expected_output = input_data.transpose((2, 0, 1))
        torch.testing.assert_close(output.numpy(), expected_output)

        input_data = torch.as_tensor(np_rng.rand(channels, height, width).astype(np.float32))
        img = transforms.ToPILImage()(input_data)  # CHW -> HWC and (* 255).byte()
        output = trans(img)  # HWC -> CHW
        expected_output = (input_data * 255).byte()
        torch.testing.assert_close(output, expected_output)

        # separate test for mode '1' PIL images
        input_data = torch.ByteTensor(1, height, width).bernoulli_()
        img = transforms.ToPILImage()(input_data.mul(255)).convert("1")
        output = trans(img).view(torch.uint8).bool().to(torch.uint8)
        torch.testing.assert_close(input_data, output)
//  [python]
def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        img = transform(img_data)
        assert img.mode == expected_mode
        if np.issubdtype(img_data.dtype, np.floating):
            img_data = (img_data * 255).astype(np.uint8)
        # note: we explicitly convert img's dtype because pytorch doesn't support uint16
        # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array
        torch.testing.assert_allclose(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype))
//  [python]
def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        img = transform(img_data)
        assert img.mode == expected_mode
        if np.issubdtype(img_data.dtype, np.floating):
            img_data = (img_data * 255).astype(np.uint8)
        # note: we explicitly convert img's dtype because pytorch doesn't support uint16
        # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array
        torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype))
//  [python]
def test_2_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "LA"  # default should assume LA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(2):
            torch.testing.assert_allclose(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_2_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "LA"  # default should assume LA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(2):
            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        to_tensor = transforms.ToTensor()

        img = transform(img_data)
        assert img.mode == expected_mode
        torch.testing.assert_allclose(expected_output, to_tensor(img).numpy()[0])
//  [python]
def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
        to_tensor = transforms.ToTensor()

        img = transform(img_data)
        assert img.mode == expected_mode
        torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0])
//  [python]
def test_2_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(2, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)
        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "LA"  # default should assume LA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode

        split = img.split()
        for i in range(2):
            torch.testing.assert_allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
//  [python]
def test_2_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(2, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)
        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "LA"  # default should assume LA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode

        split = img.split()
        for i in range(2):
            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
//  [python]
def test_3_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGB"  # default should assume RGB
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(3):
            torch.testing.assert_allclose(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_3_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGB"  # default should assume RGB
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(3):
            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_3_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(3, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGB"  # default should assume RGB
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(3):
            torch.testing.assert_allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())