jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_replace_linear(self):
        device = torch.device("cuda")
        dtype = torch.float32
        batch, head_dim, head_num, hidden = 128, 64, 8, 768
        x = torch.randn(batch, head_dim, head_num, hidden, device=device, dtype=dtype)
        layer = MyModule(True, hidden, hidden)
        layer.reset_parameters()
        replace_linear(layer, "")
        layer.to(device)
        self.assertTrue(isinstance(layer.fc1, Linear2D))
        layer_fused = MyModule(True, hidden, hidden).to(device)
        layer_fused.load_state_dict(layer.state_dict())
        with torch.cuda.amp.autocast(dtype=dtype):
            y = layer(x)
            y1 = layer_fused(x)
            torch.testing.assert_close(y, y1, rtol=1e-4, atol=1e-4)
            Y = torch.randn_like(y)
            loss = torch.nn.functional.mse_loss(y, Y)
            loss1 = torch.nn.functional.mse_loss(y1, Y)
            loss.backward()
            loss1.backward()
            name2grad = {}
            for name, p in layer.named_parameters():
                name2grad[name] = p.grad
            for name, p in layer_fused.named_parameters():
                torch.testing.assert_close(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
//  [python]
def verify_step(model, args, jit_script=True, with_backward=False, tol=1e-5):
    """Verify the results between CPU and Lazy"""
    if with_backward:
        args_ratex = []
        args_cpu = args
        for x in args:
            n_x = np.copy(x.data)
            t_x_ratex = torch.tensor(n_x, device="lazy", dtype=torch.float32, requires_grad=True)
            args_ratex.append(t_x_ratex)
        out_cpu, grad_cpu = run_step("cpu", model, args_cpu, with_backward=with_backward)
        out_lazy, grad_lazy = run_step("lazy", model, args_ratex, jit_script, with_backward)
        t.assert_close(grad_cpu, grad_lazy, rtol=tol, atol=tol)
        t.assert_close(out_cpu, out_lazy, rtol=tol, atol=tol)
    else:
        t.assert_close(
            run_step("cpu", model, args, with_backward=with_backward),
            run_step("lazy", model, args, jit_script=jit_script, with_backward=with_backward),
            rtol=tol,
            atol=tol,
        )
//  [python]
def verify_step(model, args, jit_script=True, with_backward=False, tol=1e-5):
    """Verify the results between CPU and Lazy"""
    if with_backward:
        args_ratex = []
        args_cpu = args
        for x in args:
            n_x = np.copy(x.data)
            t_x_ratex = torch.tensor(n_x, device="lazy", dtype=torch.float32, requires_grad=True)
            args_ratex.append(t_x_ratex)
        out_cpu, grad_cpu = run_step("cpu", model, args_cpu, with_backward=with_backward)
        out_lazy, grad_lazy = run_step("lazy", model, args_ratex, jit_script, with_backward)
        torch.testing.assert_close(grad_cpu, grad_lazy, rtol=tol, atol=tol)
        torch.testing.assert_close(out_cpu, out_lazy, rtol=tol, atol=tol)
    else:
        torch.testing.assert_close(
            run_step("cpu", model, args, with_backward=with_backward),
            run_step("lazy", model, args, jit_script=jit_script, with_backward=with_backward),
            rtol=tol,
            atol=tol,
        )
//  [python]
def test_word_error_rate_with_valid_input(self) -> None:
        t.assert_allclose(
            word_error_rate("hello meta", "hello metaverse"),
            torch.tensor(0.5, dtype=torch.float64),
        )
        t.assert_allclose(
            word_error_rate("hello meta", "hello meta"),
            torch.tensor(0.0, dtype=torch.float64),
        )
        t.assert_allclose(
            word_error_rate("this is the prediction", "this is the reference"),
            torch.tensor(0.25, dtype=torch.float64),
        )
        t.assert_allclose(
            word_error_rate(
                ["hello world", "welcome to the facebook"],
                ["hello metaverse", "welcome to meta"],
            ),
            torch.tensor(0.6, dtype=torch.float64),
        )
        t.assert_allclose(
            word_error_rate(
                [
                    "hello metaverse",
                    "come to the facebook",
                    "this is reference",
                    "there is the other one",
                ],
                [
                    "hello world",
                    "welcome to meta",
                    "this is reference",
                    "there is another one",
                ],
            ),
            torch.tensor(0.5, dtype=torch.float64),
        )
//  [python]
def test_word_error_rate_with_valid_input(self) -> None:
        torch.testing.assert_close(
            word_error_rate("hello meta", "hello metaverse"),
            torch.tensor(0.5, dtype=torch.float64),
        )
        torch.testing.assert_close(
            word_error_rate("hello meta", "hello meta"),
            torch.tensor(0.0, dtype=torch.float64),
        )
        torch.testing.assert_close(
            word_error_rate("this is the prediction", "this is the reference"),
            torch.tensor(0.25, dtype=torch.float64),
        )
        torch.testing.assert_close(
            word_error_rate(
                ["hello world", "welcome to the facebook"],
                ["hello metaverse", "welcome to meta"],
            ),
            torch.tensor(0.6, dtype=torch.float64),
        )
        torch.testing.assert_close(
            word_error_rate(
                [
                    "hello metaverse",
                    "come to the facebook",
                    "this is reference",
                    "there is the other one",
                ],
                [
                    "hello world",
                    "welcome to meta",
                    "this is reference",
                    "there is another one",
                ],
            ),
            torch.tensor(0.5, dtype=torch.float64),
        )
//  [python]
def test_hf_loading_full_model_equivalence(tmp_path: Path, tiny_bert_model, tiny_bert_tokenizer):
    pytest.importorskip('transformers')

    trainer1 = get_lm_trainer(tiny_bert_model, tiny_bert_tokenizer, str(tmp_path))
    trainer1.fit()

    hf_loaded_model, hf_loaded_tokenizer = HuggingFaceModel.hf_from_composer_checkpoint(
        checkpoint_path=str(tmp_path / 'hf-checkpoint.pt'),
    )

    trainer2 = get_lm_trainer(
        hf_loaded_model,
        hf_loaded_tokenizer,
        str(tmp_path),
        load_path=str(tmp_path / 'hf-checkpoint.pt'),
    )

    # loading from the last checkpoint gets you the same model
    for p1, p2 in zip(trainer1.state.model.parameters(), trainer2.state.model.parameters()):
        torch.allclose(p1, p2)
//  [python]
def test_hf_loading_full_model_equivalence(tmp_path: Path, tiny_bert_model, tiny_bert_tokenizer):
    pytest.importorskip('transformers')

    trainer1 = get_lm_trainer(tiny_bert_model, tiny_bert_tokenizer, str(tmp_path))
    trainer1.fit()

    hf_loaded_model, hf_loaded_tokenizer = HuggingFaceModel.hf_from_composer_checkpoint(
        checkpoint_path=str(tmp_path / 'hf-checkpoint.pt'),
    )

    trainer2 = get_lm_trainer(
        hf_loaded_model,
        hf_loaded_tokenizer,
        str(tmp_path),
        load_path=str(tmp_path / 'hf-checkpoint.pt'),
    )

    # loading from the last checkpoint gets you the same model
    for p1, p2 in zip(trainer1.state.model.parameters(), trainer2.state.model.parameters()):
        torch.testing.assert_close(p1, p2)
//  [python]
def _check_kernel_cuda_vs_cpu(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces closes results for inputs on GPU and CPU."""
    if input.device.type != "cuda":
        return

    input_cuda = input.as_subclass(torch.Tensor)
    input_cpu = input_cuda.to("cpu")

    with freeze_rng_state():
        actual = kernel(input_cuda, *args, **kwargs)
    with freeze_rng_state():
        expected = kernel(input_cpu, *args, **kwargs)

    t.assert_close(actual, expected, check_device=False, rtol=rtol, atol=atol)
//  [python]
def _check_kernel_cuda_vs_cpu(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces closes results for inputs on GPU and CPU."""
    if input.device.type != "cuda":
        return

    input_cuda = input.as_subclass(torch.Tensor)
    input_cpu = input_cuda.to("cpu")

    with freeze_rng_state():
        actual = kernel(input_cuda, *args, **kwargs)
    with freeze_rng_state():
        expected = kernel(input_cpu, *args, **kwargs)

    assert_close(actual, expected, check_device=False, rtol=rtol, atol=atol)
//  [python]
def check_grad(solve_func, inputs, eps, atol, rtol):
    assert gradcheck(solve_func, inputs, eps=eps, atol=atol)

    A_val, b = inputs[0], inputs[1]
    # Check that the gradient works correctly for floating point data
    out = solve_func(*inputs).sum()
    gA, gb = grad(out, (A_val, b))

    A_float = A_val.float()
    b_float = b.float()
    inputs2 = (A_float, b_float) + inputs[2:]
    out_float = solve_func(*inputs2).sum()
    gA_float, gb_float = grad(out_float, (A_float, b_float))

    # This mostly checks that backward() is not accumulating
    # additional floating point errors on top of those expected by
    # converting the input from double to float. Naive float casting
    # in the backward python ops can result in differences in order of magnitude
    # even for well-conditioned systems. These checks cover that case.
    #
    # Do note that, in general, it's possible to construct very
    # ill-conditioned systems where the initial loss or precision is enough
    # to get large gradient errors no matter what we do. This checks
    # are not meant to capture such scenarios.
    t.assert_close(gA, gA_float.double(), rtol=rtol, atol=atol)
    t.assert_close(gb, gb_float.double(), rtol=rtol, atol=atol)
//  [python]
def check_grad(solve_func, inputs, eps, atol, rtol):
    assert gradcheck(solve_func, inputs, eps=eps, atol=atol)

    A_val, b = inputs[0], inputs[1]
    # Check that the gradient works correctly for floating point data
    out = solve_func(*inputs).sum()
    gA, gb = grad(out, (A_val, b))

    A_float = A_val.float()
    b_float = b.float()
    inputs2 = (A_float, b_float) + inputs[2:]
    out_float = solve_func(*inputs2).sum()
    gA_float, gb_float = grad(out_float, (A_float, b_float))

    # This mostly checks that backward() is not accumulating
    # additional floating point errors on top of those expected by
    # converting the input from double to float. Naive float casting
    # in the backward python ops can result in differences in order of magnitude
    # even for well-conditioned systems. These checks cover that case.
    #
    # Do note that, in general, it's possible to construct very
    # ill-conditioned systems where the initial loss or precision is enough
    # to get large gradient errors no matter what we do. This checks
    # are not meant to capture such scenarios.
    torch.testing.assert_close(gA, gA_float.double(), rtol=rtol, atol=atol)
    torch.testing.assert_close(gb, gb_float.double(), rtol=rtol, atol=atol)
//  [python]
def test_parallel_embedding(self):
        torch.manual_seed(9999)
        input_tensor = torch.randint(low=0, high=NUM_EMBEDDINGS, size=(INPUT_SIZE,)).cuda()
        golden_model = VocabEmbedding().cuda()
        golden_out = golden_model(input_tensor)
        loss = golden_out.mean()
        loss.backward()
        golden_weight_grad = golden_model.emb.weight.grad

        device_mesh = DeviceMesh("cuda", list(range(self.world_size)))
        torch.manual_seed(9999)
        input_tensor = torch.randint(low=0, high=NUM_EMBEDDINGS, size=(INPUT_SIZE,)).cuda()
        dmodel = VocabEmbedding()
        param_sharding_plan = {"emb.weight": [Shard(0)]}
        fwd_resharding_plan = {".input": [[Replicate()]]}
        parallelize_module(dmodel, device_mesh, {"parameter": param_sharding_plan, "forward": fwd_resharding_plan})
        dt_out = dmodel(input_tensor)
        local_out = dt_out.redistribute(placements=[Replicate()] * device_mesh.ndim).to_local()
        loss = local_out.mean()
        loss.backward()
        self.assertEqual(len(dmodel.emb.weight.grad.placements), 1)
        self.assertTrue(dmodel.emb.weight.grad.placements[0].is_shard(0))
        dmodel_weight_grad = dmodel.emb.weight.grad.redistribute(placements=[Replicate()] * device_mesh.ndim).to_local()

        self.assertEqual(golden_out, local_out)

        self.assertEqual(golden_weight_grad, dmodel_weight_grad)
//  [python]
def test_parallel_embedding(self):
        torch.manual_seed(9999)
        input_tensor = torch.randint(low=0, high=NUM_EMBEDDINGS, size=(INPUT_SIZE,)).cuda()
        golden_model = VocabEmbedding().cuda()
        golden_out = golden_model(input_tensor)
        loss = golden_out.mean()
        loss.backward()
        golden_weight_grad = golden_model.emb.weight.grad

        device_mesh = DeviceMesh("cuda", list(range(self.world_size)))
        torch.manual_seed(9999)
        input_tensor = torch.randint(low=0, high=NUM_EMBEDDINGS, size=(INPUT_SIZE,)).cuda()
        dmodel = VocabEmbedding()
        param_sharding_plan = {"emb.weight": [Shard(0)]}
        fwd_resharding_plan = {".input": [[Replicate()]]}
        parallelize_module(dmodel, device_mesh, {"parameter": param_sharding_plan, "forward": fwd_resharding_plan})
        dt_out = dmodel(input_tensor)
        local_out = dt_out.redistribute(placements=[Replicate()] * device_mesh.ndim).to_local()
        loss = local_out.mean()
        loss.backward()
        self.assertEqual(len(dmodel.emb.weight.grad.placements), 1)
        self.assertTrue(dmodel.emb.weight.grad.placements[0].is_shard(0))
        dmodel_weight_grad = dmodel.emb.weight.grad.redistribute(placements=[Replicate()] * device_mesh.ndim).to_local()

        torch.testing.assert_close(golden_out, local_out)

        torch.testing.assert_close(golden_weight_grad, dmodel_weight_grad)
//  [python]
def test_peft_trains_and_loads(tiny_gpt2_model, tiny_gpt2_tokenizer, gpt2_peft_config, tmp_path, should_save_peft_only):
    pytest.importorskip('peft')

    trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        should_save_peft_only=should_save_peft_only,
    )
    trainer.fit()

    load_trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        load_path=str(tmp_path / 'hf-checkpoint.pt'),
        should_save_peft_only=should_save_peft_only,
    )

    for p1, p2 in zip(trainer.state.model.parameters(), load_trainer.state.model.parameters()):
        torch.testing.assert_allclose(p1, p2)
//  [python]
def test_peft_trains_and_loads(tiny_gpt2_model, tiny_gpt2_tokenizer, gpt2_peft_config, tmp_path, should_save_peft_only):
    pytest.importorskip('peft')

    trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        should_save_peft_only=should_save_peft_only,
    )
    trainer.fit()

    load_trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        load_path=str(tmp_path / 'hf-checkpoint.pt'),
        should_save_peft_only=should_save_peft_only,
    )

    for p1, p2 in zip(trainer.state.model.parameters(), load_trainer.state.model.parameters()):
        torch.testing.assert_close(p1, p2)
//  [python]
def _check_kernel_scripted_vs_eager(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel is scriptable and if the scripted output is close to the eager one."""
    if input.device.type != "cpu":
        return

    kernel_scripted = _script(kernel)

    input = input.as_subclass(torch.Tensor)
    with ignore_jit_no_profile_information_warning():
        actual = kernel_scripted(input, *args, **kwargs)
    expected = kernel(input, *args, **kwargs)

    t.assert_close(actual, expected, rtol=rtol, atol=atol)
//  [python]
def _check_kernel_scripted_vs_eager(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel is scriptable and if the scripted output is close to the eager one."""
    if input.device.type != "cpu":
        return

    kernel_scripted = _script(kernel)

    input = input.as_subclass(torch.Tensor)
    with ignore_jit_no_profile_information_warning():
        actual = kernel_scripted(input, *args, **kwargs)
    expected = kernel(input, *args, **kwargs)

    assert_close(actual, expected, rtol=rtol, atol=atol)
//  [python]
def test_image_correctness(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.horizontal_flip(F.to_pil_image(image)))

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_image_correctness(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.horizontal_flip(F.to_pil_image(image)))

        torch.testing.assert_close(actual, expected)
//  [python]
def test_peft_write_hf_from_composer(
    tiny_gpt2_model,
    tiny_gpt2_tokenizer,
    gpt2_peft_config,
    tmp_path,
    should_save_peft_only,
):
    peft = pytest.importorskip('peft')
    transformers = pytest.importorskip('transformers')

    # Simulate a local model instead of a hub model
    tiny_gpt2_model.save_pretrained(tmp_path / 'hf-save-to-load')
    tiny_gpt2_model = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / 'hf-save-to-load')

    trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        should_save_peft_only=should_save_peft_only,
    )
    trainer.fit()

    from composer.models.huggingface import write_huggingface_pretrained_from_composer_checkpoint
    write_huggingface_pretrained_from_composer_checkpoint(
        str(tmp_path / 'hf-checkpoint.pt'),
        tmp_path / 'hf-save-pretrained',
    )

    # Test we can load back in using transformers interface
    loaded_hf_model = transformers.AutoModelForCausalLM.from_pretrained(str(tmp_path / 'hf-save-pretrained'))
    for p1, p2 in zip(trainer.state.model.model.parameters(), loaded_hf_model.parameters()):
        torch.allclose(p1, p2)

    # Test we can load back in using peft interface
    loaded_peft_model = peft.PeftModelForCausalLM.from_pretrained(tiny_gpt2_model, str(tmp_path / 'hf-save-pretrained'))
    for p1, p2 in zip(trainer.state.model.model.parameters(), loaded_peft_model.parameters()):
        torch.allclose(p1, p2)
//  [python]
def test_peft_write_hf_from_composer(
    tiny_gpt2_model,
    tiny_gpt2_tokenizer,
    gpt2_peft_config,
    tmp_path,
    should_save_peft_only,
):
    peft = pytest.importorskip('peft')
    transformers = pytest.importorskip('transformers')

    # Simulate a local model instead of a hub model
    tiny_gpt2_model.save_pretrained(tmp_path / 'hf-save-to-load')
    tiny_gpt2_model = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / 'hf-save-to-load')

    trainer = get_lm_trainer(
        tiny_gpt2_model,
        tiny_gpt2_tokenizer,
        str(tmp_path),
        peft_config=gpt2_peft_config,
        device_train_microbatch_size=1,
        mlm=False,
        should_save_peft_only=should_save_peft_only,
    )
    trainer.fit()

    from composer.models.huggingface import write_huggingface_pretrained_from_composer_checkpoint
    write_huggingface_pretrained_from_composer_checkpoint(
        str(tmp_path / 'hf-checkpoint.pt'),
        tmp_path / 'hf-save-pretrained',
    )

    # Test we can load back in using transformers interface
    loaded_hf_model = transformers.AutoModelForCausalLM.from_pretrained(str(tmp_path / 'hf-save-pretrained'))
    for p1, p2 in zip(trainer.state.model.model.parameters(), loaded_hf_model.parameters()):
        torch.testing.assert_close(p1, p2)

    # Test we can load back in using peft interface
    loaded_peft_model = peft.PeftModelForCausalLM.from_pretrained(tiny_gpt2_model, str(tmp_path / 'hf-save-pretrained'))
    for p1, p2 in zip(trainer.state.model.model.parameters(), loaded_peft_model.parameters()):
        torch.testing.assert_close(p1, p2)
//  [python]
def test_bounding_boxes_correctness(self, format, fn):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = fn(bounding_boxes)
        expected = self._reference_horizontal_flip_bounding_boxes(bounding_boxes)

        t.assert_close(actual, expected)
//  [python]
def test_bounding_boxes_correctness(self, format, fn):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = fn(bounding_boxes)
        expected = self._reference_horizontal_flip_bounding_boxes(bounding_boxes)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_image_correctness(self, size, interpolation, use_max_size, fn):
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)

        actual = fn(image, size=size, interpolation=interpolation, **max_size_kwarg, antialias=True)
        expected = F.to_image(F.resize(F.to_pil_image(image), size=size, interpolation=interpolation, **max_size_kwarg))

        self._check_output_size(image, actual, size=size, **max_size_kwarg)
        t.assert_close(actual, expected, atol=1, rtol=0)
//  [python]
def test_image_correctness(self, size, interpolation, use_max_size, fn):
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)

        actual = fn(image, size=size, interpolation=interpolation, **max_size_kwarg, antialias=True)
        expected = F.to_image(F.resize(F.to_pil_image(image), size=size, interpolation=interpolation, **max_size_kwarg))

        self._check_output_size(image, actual, size=size, **max_size_kwarg)
        torch.testing.assert_close(actual, expected, atol=1, rtol=0)
//  [python]
def test_bounding_boxes_correctness(self, format, size, use_max_size, fn):
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

        bounding_boxes = make_bounding_boxes(format=format, canvas_size=self.INPUT_SIZE)

        actual = fn(bounding_boxes, size=size, **max_size_kwarg)
        expected = self._reference_resize_bounding_boxes(bounding_boxes, size=size, **max_size_kwarg)

        self._check_output_size(bounding_boxes, actual, size=size, **max_size_kwarg)
        t.assert_close(actual, expected)
//  [python]
def test_bounding_boxes_correctness(self, format, size, use_max_size, fn):
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

        bounding_boxes = make_bounding_boxes(format=format, canvas_size=self.INPUT_SIZE)

        actual = fn(bounding_boxes, size=size, **max_size_kwarg)
        expected = self._reference_resize_bounding_boxes(bounding_boxes, size=size, **max_size_kwarg)

        self._check_output_size(bounding_boxes, actual, size=size, **max_size_kwarg)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_functional_bounding_boxes_correctness(self, format, angle, translate, scale, shear, center):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = F.affine(
            bounding_boxes,
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )
        expected = self._reference_affine_bounding_boxes(
            bounding_boxes,
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )

        t.assert_close(actual, expected)
//  [python]
def test_functional_bounding_boxes_correctness(self, format, angle, translate, scale, shear, center):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = F.affine(
            bounding_boxes,
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )
        expected = self._reference_affine_bounding_boxes(
            bounding_boxes,
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )

        torch.testing.assert_close(actual, expected)
//  [python]
def _check_kernel_batched_vs_unbatched(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces close results for batched and unbatched inputs."""
    unbatched_input = input.as_subclass(torch.Tensor)

    for batch_dims in [(2,), (2, 1)]:
        repeats = [*batch_dims, *[1] * input.ndim]

        actual = kernel(unbatched_input.repeat(repeats), *args, **kwargs)

        expected = kernel(unbatched_input, *args, **kwargs)
        # We can't directly call `.repeat()` on the output, since some kernel also return some additional metadata
        if isinstance(expected, torch.Tensor):
            expected = expected.repeat(repeats)
        else:
            tensor, *metadata = expected
            expected = (tensor.repeat(repeats), *metadata)

        t.assert_close(actual, expected, rtol=rtol, atol=atol)

    for degenerate_batch_dims in [(0,), (5, 0), (0, 5)]:
        degenerate_batched_input = torch.empty(
            degenerate_batch_dims + input.shape, dtype=input.dtype, device=input.device
        )

        output = kernel(degenerate_batched_input, *args, **kwargs)
        # Most kernels just return a tensor, but some also return some additional metadata
        if not isinstance(output, torch.Tensor):
            output, *_ = output

        assert output.shape[: -input.ndim] == degenerate_batch_dims
//  [python]
def _check_kernel_batched_vs_unbatched(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces close results for batched and unbatched inputs."""
    unbatched_input = input.as_subclass(torch.Tensor)

    for batch_dims in [(2,), (2, 1)]:
        repeats = [*batch_dims, *[1] * input.ndim]

        actual = kernel(unbatched_input.repeat(repeats), *args, **kwargs)

        expected = kernel(unbatched_input, *args, **kwargs)
        # We can't directly call `.repeat()` on the output, since some kernel also return some additional metadata
        if isinstance(expected, torch.Tensor):
            expected = expected.repeat(repeats)
        else:
            tensor, *metadata = expected
            expected = (tensor.repeat(repeats), *metadata)

        assert_close(actual, expected, rtol=rtol, atol=atol)

    for degenerate_batch_dims in [(0,), (5, 0), (0, 5)]:
        degenerate_batched_input = torch.empty(
            degenerate_batch_dims + input.shape, dtype=input.dtype, device=input.device
        )

        output = kernel(degenerate_batched_input, *args, **kwargs)
        # Most kernels just return a tensor, but some also return some additional metadata
        if not isinstance(output, torch.Tensor):
            output, *_ = output

        assert output.shape[: -input.ndim] == degenerate_batch_dims
//  [python]
def _check_transform_v1_compatibility(transform, input, *, rtol, atol):
    """If the transform defines the ``_v1_transform_cls`` attribute, checks if the transform has a public, static
    ``get_params`` method that is the v1 equivalent, the output is close to v1, is scriptable, and the scripted version
    can be called without error."""
    if not (type(input) is torch.Tensor or isinstance(input, PIL.Image.Image)):
        return

    v1_transform_cls = transform._v1_transform_cls
    if v1_transform_cls is None:
        return

    if hasattr(v1_transform_cls, "get_params"):
        assert type(transform).get_params is v1_transform_cls.get_params

    v1_transform = v1_transform_cls(**transform._extract_params_for_v1_transform())

    with freeze_rng_state():
        output_v2 = transform(input)

    with freeze_rng_state():
        output_v1 = v1_transform(input)

    t.assert_close(F.to_image(output_v2), F.to_image(output_v1), rtol=rtol, atol=atol)

    if isinstance(input, PIL.Image.Image):
        return

    _script(v1_transform)(input)
//  [python]
def _check_transform_v1_compatibility(transform, input, *, rtol, atol):
    """If the transform defines the ``_v1_transform_cls`` attribute, checks if the transform has a public, static
    ``get_params`` method that is the v1 equivalent, the output is close to v1, is scriptable, and the scripted version
    can be called without error."""
    if not (type(input) is torch.Tensor or isinstance(input, PIL.Image.Image)):
        return

    v1_transform_cls = transform._v1_transform_cls
    if v1_transform_cls is None:
        return

    if hasattr(v1_transform_cls, "get_params"):
        assert type(transform).get_params is v1_transform_cls.get_params

    v1_transform = v1_transform_cls(**transform._extract_params_for_v1_transform())

    with freeze_rng_state():
        output_v2 = transform(input)

    with freeze_rng_state():
        output_v1 = v1_transform(input)

    assert_close(F.to_image(output_v2), F.to_image(output_v1), rtol=rtol, atol=atol)

    if isinstance(input, PIL.Image.Image):
        return

    _script(v1_transform)(input)
//  [python]
def test_image_correctness(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.vertical_flip(F.to_pil_image(image)))

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_image_correctness(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.vertical_flip(F.to_pil_image(image)))

        torch.testing.assert_close(actual, expected)
//  [python]
def test_bounding_boxes_correctness(self, format, fn):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = fn(bounding_boxes)
        expected = self._reference_vertical_flip_bounding_boxes(bounding_boxes)

        t.assert_close(actual, expected)
//  [python]
def test_bounding_boxes_correctness(self, format, fn):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = fn(bounding_boxes)
        expected = self._reference_vertical_flip_bounding_boxes(bounding_boxes)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_transform_bounding_boxes_correctness(self, format, center, seed):
        bounding_boxes = make_bounding_boxes(format=format)

        transform = transforms.RandomAffine(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center)

        torch.manual_seed(seed)
        params = transform._get_params([bounding_boxes])

        torch.manual_seed(seed)
        actual = transform(bounding_boxes)

        expected = self._reference_affine_bounding_boxes(bounding_boxes, **params, center=center)

        t.assert_close(actual, expected)
//  [python]
def test_transform_bounding_boxes_correctness(self, format, center, seed):
        bounding_boxes = make_bounding_boxes(format=format)

        transform = transforms.RandomAffine(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center)

        torch.manual_seed(seed)
        params = transform._get_params([bounding_boxes])

        torch.manual_seed(seed)
        actual = transform(bounding_boxes)

        expected = self._reference_affine_bounding_boxes(bounding_boxes, **params, center=center)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_functional_bounding_boxes_correctness(self, format, angle, expand, center):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = F.rotate(bounding_boxes, angle=angle, expand=expand, center=center)
        expected = self._reference_rotate_bounding_boxes(bounding_boxes, angle=angle, expand=expand, center=center)

        t.assert_close(actual, expected)
        t.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
//  [python]
def test_functional_bounding_boxes_correctness(self, format, angle, expand, center):
        bounding_boxes = make_bounding_boxes(format=format)

        actual = F.rotate(bounding_boxes, angle=angle, expand=expand, center=center)
        expected = self._reference_rotate_bounding_boxes(bounding_boxes, angle=angle, expand=expand, center=center)

        torch.testing.assert_close(actual, expected)
        torch.testing.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
//  [python]
def test_functional_image_fast_path_correctness(self, size, angle, expand):
        image = make_image(size, dtype=torch.uint8, device="cpu")

        actual = F.rotate(image, angle=angle, expand=expand)
        expected = F.to_image(F.rotate(F.to_pil_image(image), angle=angle, expand=expand))

        t.assert_close(actual, expected)
//  [python]
def test_functional_image_fast_path_correctness(self, size, angle, expand):
        image = make_image(size, dtype=torch.uint8, device="cpu")

        actual = F.rotate(image, angle=angle, expand=expand)
        expected = F.to_image(F.rotate(F.to_pil_image(image), angle=angle, expand=expand))

        torch.testing.assert_close(actual, expected)
//  [python]
def test_image_correctness(self, brightness_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_brightness(image, brightness_factor=brightness_factor)
        expected = F.to_image(F.adjust_brightness(F.to_pil_image(image), brightness_factor=brightness_factor))

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_image_correctness(self, brightness_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_brightness(image, brightness_factor=brightness_factor)
        expected = F.to_image(F.adjust_brightness(F.to_pil_image(image), brightness_factor=brightness_factor))

        torch.testing.assert_close(actual, expected)
//  [python]
def test_transform_bounding_boxes_correctness(self, format, expand, center, seed):
        bounding_boxes = make_bounding_boxes(format=format)

        transform = transforms.RandomRotation(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, expand=expand, center=center)

        torch.manual_seed(seed)
        params = transform._get_params([bounding_boxes])

        torch.manual_seed(seed)
        actual = transform(bounding_boxes)

        expected = self._reference_rotate_bounding_boxes(bounding_boxes, **params, expand=expand, center=center)

        self.assertEqual(actual.shape, expected.shape)
        self.assertTrue(torch.allclose(actual, expected))
        self.assertTrue(torch.allclose(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0))
//  [python]
def test_transform_bounding_boxes_correctness(self, format, expand, center, seed):
        bounding_boxes = make_bounding_boxes(format=format)

        transform = transforms.RandomRotation(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, expand=expand, center=center)

        torch.manual_seed(seed)
        params = transform._get_params([bounding_boxes])

        torch.manual_seed(seed)
        actual = transform(bounding_boxes)

        expected = self._reference_rotate_bounding_boxes(bounding_boxes, **params, expand=expand, center=center)

        torch.testing.assert_close(actual, expected)
        torch.testing.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
//  [python]
def test_image_correctness(self, input_dtype, output_dtype, device, scale):
        if input_dtype.is_floating_point and output_dtype == torch.int64:
            pytest.xfail("float to int64 conversion is not supported")

        input = make_image(dtype=input_dtype, device=device)

        out = F.to_dtype(input, dtype=output_dtype, scale=scale)
        expected = self.reference_convert_dtype_image_tensor(input, dtype=output_dtype, scale=scale)

        if input_dtype.is_floating_point and not output_dtype.is_floating_point and scale:
            torch.testing.assert_allclose(out, expected, atol=1, rtol=0)
        else:
            torch.testing.assert_allclose(out, expected)
//  [python]
def test_image_correctness(self, input_dtype, output_dtype, device, scale):
        if input_dtype.is_floating_point and output_dtype == torch.int64:
            pytest.xfail("float to int64 conversion is not supported")

        input = make_image(dtype=input_dtype, device=device)

        out = F.to_dtype(input, dtype=output_dtype, scale=scale)
        expected = self.reference_convert_dtype_image_tensor(input, dtype=output_dtype, scale=scale)

        if input_dtype.is_floating_point and not output_dtype.is_floating_point and scale:
            torch.testing.assert_close(out, expected, atol=1, rtol=0)
        else:
            torch.testing.assert_close(out, expected)
//  [python]
def test_image_correctness(self, permutation, batch_dims):
        image = make_image(batch_dims=batch_dims)

        actual = F.permute_channels(image, permutation=permutation)
        expected = self.reference_image_correctness(image, permutation=permutation)

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_image_correctness(self, permutation, batch_dims):
        image = make_image(batch_dims=batch_dims)

        actual = F.permute_channels(image, permutation=permutation)
        expected = self.reference_image_correctness(image, permutation=permutation)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_correctness_image(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.autocontrast(F.to_pil_image(image)))

        t.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_image(self, fn):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = fn(image)
        expected = F.to_image(F.autocontrast(F.to_pil_image(image)))

        assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_functional_image_correctness(self, interpolation):
        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)

        actual = F.resized_crop(
            image, **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation, antialias=True
        )
        expected = F.to_image(
            F.resized_crop(
                F.to_pil_image(image), **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation
            )
        )

        torch.testing.assert_allclose(actual, expected, atol=1, rtol=0)
//  [python]
def test_functional_image_correctness(self, interpolation):
        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)

        actual = F.resized_crop(
            image, **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation, antialias=True
        )
        expected = F.to_image(
            F.resized_crop(
                F.to_pil_image(image), **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation
            )
        )

        torch.testing.assert_close(actual, expected, atol=1, rtol=0)
//  [python]
def test_correctness_perspective_bounding_boxes(self, startpoints, endpoints, format, dtype, device):
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)

        actual = F.perspective(bounding_boxes, startpoints=startpoints, endpoints=endpoints)
        expected = self._reference_perspective_bounding_boxes(
            bounding_boxes, startpoints=startpoints, endpoints=endpoints
        )

        t.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_perspective_bounding_boxes(self, startpoints, endpoints, format, dtype, device):
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)

        actual = F.perspective(bounding_boxes, startpoints=startpoints, endpoints=endpoints)
        expected = self._reference_perspective_bounding_boxes(
            bounding_boxes, startpoints=startpoints, endpoints=endpoints
        )

        assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_image(self, contrast_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_contrast(image, contrast_factor=contrast_factor)
        expected = F.to_image(F.adjust_contrast(F.to_pil_image(image), contrast_factor=contrast_factor))

        torch.testing.assert_allclose(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_image(self, contrast_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_contrast(image, contrast_factor=contrast_factor)
        expected = F.to_image(F.adjust_contrast(F.to_pil_image(image), contrast_factor=contrast_factor))

        assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_functional_image_correctness(self, dimensions, kernel_size, sigma, dtype, device):
        if dtype is torch.float16 and device == "cpu":
            pytest.skip("The CPU implementation of float16 on CPU differs from opencv")

        num_channels, height, width = dimensions

        reference_results_key = f"{height}_{width}_{num_channels}__{kernel_size[0]}_{kernel_size[1]}_{sigma}"
        expected = (
            torch.tensor(self.REFERENCE_GAUSSIAN_BLUR_IMAGE_RESULTS[reference_results_key])
            .reshape(height, width, num_channels)
            .permute(2, 0, 1)
            .to(dtype=dtype, device=device)
        )

        image = tv_tensors.Image(
            torch.arange(num_channels * height * width, dtype=torch.uint8)
            .reshape(height, width, num_channels)
            .permute(2, 0, 1),
            dtype=dtype,
            device=device,
        )

        actual = F.gaussian_blur_image(image, kernel_size=kernel_size, sigma=sigma)

        torch.testing.assert_allclose(actual, expected, rtol=0, atol=1)
//  [python]
def test_functional_image_correctness(self, dimensions, kernel_size, sigma, dtype, device):
        if dtype is torch.float16 and device == "cpu":
            pytest.skip("The CPU implementation of float16 on CPU differs from opencv")

        num_channels, height, width = dimensions

        reference_results_key = f"{height}_{width}_{num_channels}__{kernel_size[0]}_{kernel_size[1]}_{sigma}"
        expected = (
            torch.tensor(self.REFERENCE_GAUSSIAN_BLUR_IMAGE_RESULTS[reference_results_key])
            .reshape(height, width, num_channels)
            .permute(2, 0, 1)
            .to(dtype=dtype, device=device)
        )

        image = tv_tensors.Image(
            torch.arange(num_channels * height * width, dtype=torch.uint8)
            .reshape(height, width, num_channels)
            .permute(2, 0, 1),
            dtype=dtype,
            device=device,
        )

        actual = F.gaussian_blur_image(image, kernel_size=kernel_size, sigma=sigma)

        torch.testing.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_image(self, saturation_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_saturation(image, saturation_factor=saturation_factor)
        expected = F.to_image(F.adjust_saturation(F.to_pil_image(image), saturation_factor=saturation_factor))

        t.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_image(self, saturation_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_saturation(image, saturation_factor=saturation_factor)
        expected = F.to_image(F.adjust_saturation(F.to_pil_image(image), saturation_factor=saturation_factor))

        assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_functional_image_correctness(self, quality):
        image = make_image()

        actual = F.jpeg(image, quality=quality)
        expected = F.to_image(F.jpeg(F.to_pil_image(image), quality=quality))

        # NOTE: this will fail if torchvision and Pillow use different JPEG encoder/decoder
        torch.testing.assert_allclose(actual, expected, rtol=0, atol=1)
//  [python]
def test_functional_image_correctness(self, quality):
        image = make_image()

        actual = F.jpeg(image, quality=quality)
        expected = F.to_image(F.jpeg(F.to_pil_image(image), quality=quality))

        # NOTE: this will fail if torchvision and Pillow use different JPEG encoder/decoder
        torch.testing.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_shear_translate(self, transform_id, magnitude, interpolation, fill, input_type):
        # ShearX/Y and TranslateX/Y are the only ops that are native to the AA transforms. They are modeled after the
        # reference implementation:
        # https://github.com/tensorflow/models/blob/885fda091c46c59d6c7bb5c7e760935eacc229da/research/autoaugment/augmentation_transforms.py#L273-L362
        # All other ops are checked in their respective dedicated tests.

        image = make_image(dtype=torch.uint8, device="cpu")
        if input_type == "PIL":
            image = F.to_pil_image(image)

        if "Translate" in transform_id:
            # For TranslateX/Y magnitude is a value in pixels
            magnitude *= min(F.get_size(image))

        actual = transforms.AutoAugment()._apply_image_or_video_transform(
            image,
            transform_id=transform_id,
            magnitude=magnitude,
            interpolation=interpolation,
            fill={type(image): fill},
        )
        expected = self._reference_shear_translate(
            image, transform_id=transform_id, magnitude=magnitude, interpolation=interpolation, fill=fill
        )

        if input_type == "PIL":
            actual, expected = F.to_image(actual), F.to_image(expected)

        if "Shear" in transform_id and input_type == "Tensor":
            mae = (actual.float() - expected.float()).abs().mean()
            assert mae < (12 if interpolation is transforms.InterpolationMode.NEAREST else 5)
        else:
            t.assert_allclose(actual, expected, rtol=0, atol=1)
//  [python]
def test_correctness_shear_translate(self, transform_id, magnitude, interpolation, fill, input_type):
        # ShearX/Y and TranslateX/Y are the only ops that are native to the AA transforms. They are modeled after the
        # reference implementation:
        # https://github.com/tensorflow/models/blob/885fda091c46c59d6c7bb5c7e760935eacc229da/research/autoaugment/augmentation_transforms.py#L273-L362
        # All other ops are checked in their respective dedicated tests.

        image = make_image(dtype=torch.uint8, device="cpu")
        if input_type == "PIL":
            image = F.to_pil_image(image)

        if "Translate" in transform_id:
            # For TranslateX/Y magnitude is a value in pixels
            magnitude *= min(F.get_size(image))

        actual = transforms.AutoAugment()._apply_image_or_video_transform(
            image,
            transform_id=transform_id,
            magnitude=magnitude,
            interpolation=interpolation,
            fill={type(image): fill},
        )
        expected = self._reference_shear_translate(
            image, transform_id=transform_id, magnitude=magnitude, interpolation=interpolation, fill=fill
        )

        if input_type == "PIL":
            actual, expected = F.to_image(actual), F.to_image(expected)

        if "Shear" in transform_id and input_type == "Tensor":
            mae = (actual.float() - expected.float()).abs().mean()
            assert mae < (12 if interpolation is transforms.InterpolationMode.NEAREST else 5)
        else:
            assert_close(actual, expected, rtol=0, atol=1)