jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_resnet_block(input_channels, output_channels, use_conv_shortcut, dropout_probability):
    block = ResNetBlock(input_channels=input_channels,
                        output_channels=output_channels,
                        use_conv_shortcut=use_conv_shortcut,
                        dropout_probability=dropout_probability,
                        zero_init_last=True)
    x = torch.randn(1, input_channels, 5, 5)
    y = block(x)
    assert y.shape == (1, output_channels, 5, 5), f'{y.shape}'
    if input_channels == output_channels:
        # Model should be exactly identity here
        torch.testing.assert_close(x, y)
//  [python]
def test_grad_scaler(self):

        torch.manual_seed(self.seed + self.rank)

        # Identical models with data-parallel and ZeRO
        ref_model, ref_optim, dist_model, dist_optim = make_models(1, 1)
        grad_scaler_args = dict(
            init_scale=3.21,
            growth_factor=1.23,
            backoff_factor=0.876,
            growth_interval=1,
        )
        ref_scaler =  torch.cuda.amp.GradScaler(**grad_scaler_args)
        dist_scaler =  torch.cuda.amp.GradScaler(**grad_scaler_args)

        # Training steps with pre-determined gradients
        xs = [3, 1, 4, 1, 5, 9]
        dys = [1, float('inf'), 1, 1, float('nan'), -1]
        for x, dy in zip(xs, dys):
            x = torch.tensor([[x]], dtype=torch.float32, device='cuda')
            dy = torch.tensor([[dy]], dtype=torch.float32, device='cuda')

            # Reference implementation
            ref_optim.zero_grad()
            y_ref = ref_model(x.detach())
            ref_scaler.scale(y_ref).backward(dy.detach())
            ref_scaler.step(ref_optim)
            ref_scaler.update()

            # Distributed implementation
            dist_optim.zero_grad()
            y_dist = dist_model(x.detach())
            dist_scaler.scale(y_dist).backward(dy.detach())
            dist_scaler.step(dist_optim)
            dist_scaler.update()

            # Check that parameters match
            for ref_param, dist_param in zip(ref_model.parameters(),
                                             dist_model.parameters()):
                torch.testing.assert_allclose(dist_param, ref_param)
//  [python]
def test_grad_scaler(self):

        torch.manual_seed(self.seed + self.rank)

        # Identical models with data-parallel and ZeRO
        ref_model, ref_optim, dist_model, dist_optim = make_models(1, 1)
        grad_scaler_args = dict(
            init_scale=3.21,
            growth_factor=1.23,
            backoff_factor=0.876,
            growth_interval=1,
        )
        ref_scaler =  torch.cuda.amp.GradScaler(**grad_scaler_args)
        dist_scaler =  torch.cuda.amp.GradScaler(**grad_scaler_args)

        # Training steps with pre-determined gradients
        xs = [3, 1, 4, 1, 5, 9]
        dys = [1, float('inf'), 1, 1, float('nan'), -1]
        for x, dy in zip(xs, dys):
            x = torch.tensor([[x]], dtype=torch.float32, device='cuda')
            dy = torch.tensor([[dy]], dtype=torch.float32, device='cuda')

            # Reference implementation
            ref_optim.zero_grad()
            y_ref = ref_model(x.detach())
            ref_scaler.scale(y_ref).backward(dy.detach())
            ref_scaler.step(ref_optim)
            ref_scaler.update()

            # Distributed implementation
            dist_optim.zero_grad()
            y_dist = dist_model(x.detach())
            dist_scaler.scale(y_dist).backward(dy.detach())
            dist_scaler.step(dist_optim)
            dist_scaler.update()

            # Check that parameters match
            for ref_param, dist_param in zip(ref_model.parameters(),
                                             dist_model.parameters()):
                torch.testing.assert_close(dist_param, ref_param)
//  [python]
def test_attention_gradient_scaling(scale):
    # Make the layer
    layer = GradientScalingLayer()
    # Need to register the backward hook to scale the gradient
    layer.register_full_backward_hook(layer.backward_hook)
    # Input vector
    x = torch.randn(1, 3, requires_grad=True)
    # Forward pass
    y = layer(x)
    # Check that the output is the same as the input
    t.assert_close(x, y)
    # Set the scale
    layer.set_scale(scale)
    # Take the gradient
    grad = torch.autograd.grad(y.sum(), x, retain_graph=True)[0]
    # Check that the gradient is 1 times the scale
    t.assert_close(grad, scale * torch.ones_like(grad))
//  [python]
def test_attention_gradient_scaling(scale):
    # Make the layer
    layer = GradientScalingLayer()
    # Need to register the backward hook to scale the gradient
    layer.register_full_backward_hook(layer.backward_hook)
    # Input vector
    x = torch.randn(1, 3, requires_grad=True)
    # Forward pass
    y = layer(x)
    # Check that the output is the same as the input
    torch.testing.assert_close(x, y)
    # Set the scale
    layer.set_scale(scale)
    # Take the gradient
    grad = torch.autograd.grad(y.sum(), x, retain_graph=True)[0]
    # Check that the gradient is 1 times the scale
    torch.testing.assert_close(grad, scale * torch.ones_like(grad))
//  [python]
def _test_precision_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: Optional[int] = None,
        average: Optional[str] = "micro",
    ) -> None:
        if num_classes is None:
            if average == "micro":
                my_compute_result = my_precision(input, target)
            else:
                my_compute_result = my_precision(input, target, average=average)
        else:
            if average == "micro":
                my_compute_result = my_precision(input, target, num_classes=num_classes)
            else:
                my_compute_result = my_precision(
                    input, target, num_classes=num_classes, average=average
                )
        target = target.squeeze().flatten()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input = input_label_ids.flatten()
        compute_result = torch.tensor(
            precision_score(target, input, average=average)
        ).to(torch.float32)
        if my_compute_result.shape != compute_result.shape:
            compute_result = torch.from_numpy(
                np.append(compute_result.numpy(), 0.0)
            ).to(torch.float32)
        t.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_precision_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: Optional[int] = None,
        average: Optional[str] = "micro",
    ) -> None:
        if num_classes is None:
            if average == "micro":
                my_compute_result = my_precision(input, target)
            else:
                my_compute_result = my_precision(input, target, average=average)
        else:
            if average == "micro":
                my_compute_result = my_precision(input, target, num_classes=num_classes)
            else:
                my_compute_result = my_precision(
                    input, target, num_classes=num_classes, average=average
                )
        target = target.squeeze().flatten()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input = input_label_ids.flatten()
        compute_result = torch.tensor(
            precision_score(target, input, average=average)
        ).to(torch.float32)
        if my_compute_result.shape != compute_result.shape:
            compute_result = torch.from_numpy(
                np.append(compute_result.numpy(), 0.0)
            ).to(torch.float32)
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multistep(n, key, device, T=11):
    torch.manual_seed(0)

    # mock data
    b = 5

    done = torch.zeros(b, T, 1, dtype=torch.bool, device=device)
    done[0, -1] = True
    done[1, -2] = True
    done[2, -3] = True
    done[3, -4] = True

    terminal = done.clone()
    terminal[:, -1] = done.sum(1) != 1

    mask = done.clone().cumsum(1).cumsum(1) >= 2
    mask = ~mask

    total_obs = torch.randn(1, T + 1, 1, device=device).expand(b, T + 1, 1)
    tensordict = TensorDict(
        source={
            key: total_obs[:, :T] * mask.to(torch.float),
            "done": done,
            "next": {
                key: total_obs[:, 1:] * mask.to(torch.float),
                "done": done,
                "reward": torch.randn(1, T, 1, device=device).expand(b, T, 1)
                * mask.to(torch.float),
            },
            "collector": {"mask": mask},
        },
        batch_size=(b, T),
    ).to(device)

    ms = MultiStep(
        0.9,
        n,
    ).to(device)
    ms_tensordict = ms(tensordict.clone())

    assert ms_tensordict.get("done").max() == 1

    if n == 1:
        assert_allclose_td(
            tensordict, ms_tensordict.select(*list(tensordict.keys(True, True)))
        )

    # assert that done at last step is similar to unterminated traj
    t.assert_close(
        ms_tensordict.get("gamma")[4], ms_tensordict.get("gamma")[0]
    )
    t.assert_close(
        ms_tensordict.get(("next", key))[4], ms_tensordict.get(("next", key))[0]
    )
    t.assert_close(
        ms_tensordict.get("steps_to_next_obs")[4],
        ms_tensordict.get("steps_to_next_obs")[0],
    )

    # check that next obs is properly replaced, or that it is terminated
    next_obs = ms_tensordict.get(key)[:, (ms.n_steps) :]
    true_next_obs = ms_tensordict.get(("next", key))[:, : -(ms.n_steps)]
    terminated = ~ms_tensordict.get("nonterminal")
    assert ((next_obs == true_next_obs).all(-1) | terminated[:, (ms.n_steps) :]).all()

    # test gamma computation
    t.assert_close(
        ms_tensordict.get("gamma"), ms.gamma ** ms_tensordict.get("steps_to_next_obs")
    )

    # test reward
    if n > 1:
        assert (
            ms_tensordict.get(("next", "reward"))
            != ms_tensordict.get(("next", "original_reward"))
        ).any()
    else:
        t.assert_close(
            ms_tensordict.get(("next", "reward")),
            ms_tensordict.get(("next", "original_reward")),
        )
//  [python]
def test_multistep(n, key, device, T=11):
    torch.manual_seed(0)

    # mock data
    b = 5

    done = torch.zeros(b, T, 1, dtype=torch.bool, device=device)
    done[0, -1] = True
    done[1, -2] = True
    done[2, -3] = True
    done[3, -4] = True

    terminal = done.clone()
    terminal[:, -1] = done.sum(1) != 1

    mask = done.clone().cumsum(1).cumsum(1) >= 2
    mask = ~mask

    total_obs = torch.randn(1, T + 1, 1, device=device).expand(b, T + 1, 1)
    tensordict = TensorDict(
        source={
            key: total_obs[:, :T] * mask.to(torch.float),
            "done": done,
            "next": {
                key: total_obs[:, 1:] * mask.to(torch.float),
                "done": done,
                "reward": torch.randn(1, T, 1, device=device).expand(b, T, 1)
                * mask.to(torch.float),
            },
            "collector": {"mask": mask},
        },
        batch_size=(b, T),
    ).to(device)

    ms = MultiStep(
        0.9,
        n,
    ).to(device)
    ms_tensordict = ms(tensordict.clone())

    assert ms_tensordict.get("done").max() == 1

    if n == 1:
        assert_allclose_td(
            tensordict, ms_tensordict.select(*list(tensordict.keys(True, True)))
        )

    # assert that done at last step is similar to unterminated traj
    torch.testing.assert_close(
        ms_tensordict.get("gamma")[4], ms_tensordict.get("gamma")[0]
    )
    torch.testing.assert_close(
        ms_tensordict.get(("next", key))[4], ms_tensordict.get(("next", key))[0]
    )
    torch.testing.assert_close(
        ms_tensordict.get("steps_to_next_obs")[4],
        ms_tensordict.get("steps_to_next_obs")[0],
    )

    # check that next obs is properly replaced, or that it is terminated
    next_obs = ms_tensordict.get(key)[:, (ms.n_steps) :]
    true_next_obs = ms_tensordict.get(("next", key))[:, : -(ms.n_steps)]
    terminated = ~ms_tensordict.get("nonterminal")
    assert ((next_obs == true_next_obs).all(-1) | terminated[:, (ms.n_steps) :]).all()

    # test gamma computation
    torch.testing.assert_close(
        ms_tensordict.get("gamma"), ms.gamma ** ms_tensordict.get("steps_to_next_obs")
    )

    # test reward
    if n > 1:
        assert (
            ms_tensordict.get(("next", "reward"))
            != ms_tensordict.get(("next", "original_reward"))
        ).any()
    else:
        torch.testing.assert_close(
            ms_tensordict.get(("next", "reward")),
            ms_tensordict.get(("next", "original_reward")),
        )
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto")
        input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
        with torch.no_grad():
            out = model(input_ids).logits.cpu()
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.5548, -2.5737, -3.0600, -2.5906, -2.8478, -2.8118, -2.9325, -2.7694]])
        torch.testing.assert_allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        EXPECTED_SLICE = torch.tensor([-5.8781, -5.8616, -0.1052, -4.7200, -5.8781, -5.8774, -5.8773, -5.8777, -5.8781, -5.8780, -5.8781, -5.8779, -1.0787,  1.7583, -5.8779, -5.8780, -5.8783, -5.8778, -5.8776, -5.8781, -5.8784, -5.8778, -5.8778, -5.8777, -5.8779, -5.8778, -5.8776, -5.8780, -5.8779, -5.8781])
        # fmt: on
        print(out[0, 0, :30])
        torch.testing.assert_allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-4, rtol=1e-4)

        del model
        backend_empty_cache(torch_device)
        gc.collect()
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto")
        input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
        with torch.no_grad():
            out = model(input_ids).logits.cpu()
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.5548, -2.5737, -3.0600, -2.5906, -2.8478, -2.8118, -2.9325, -2.7694]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        EXPECTED_SLICE = torch.tensor([-5.8781, -5.8616, -0.1052, -4.7200, -5.8781, -5.8774, -5.8773, -5.8777, -5.8781, -5.8780, -5.8781, -5.8779, -1.0787,  1.7583, -5.8779, -5.8780, -5.8783, -5.8778, -5.8776, -5.8781, -5.8784, -5.8778, -5.8778, -5.8777, -5.8779, -5.8778, -5.8776, -5.8780, -5.8779, -5.8781])
        # fmt: on
        print(out[0, 0, :30])
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-4, rtol=1e-4)

        del model
        backend_empty_cache(torch_device)
        gc.collect()
//  [python]
def test_mbgmv(self, input, paged_lora_a, paged_lora_b, out_head_size,
                   adapter_ids, scaling, page_table, ranks, page_start, gt):
        max_rank = page_table.size(-1)

        xa = mbgmv_a(input,
                     paged_lora_a,
                     adapter_ids=adapter_ids,
                     rank_page_table=page_table,
                     rank_page_start=page_start,
                     ranks=ranks,
                     max_rank=max_rank)

        output = mbgmv_b(xa,
                         paged_lora_b[..., :out_head_size],
                         adapter_ids=adapter_ids,
                         scaling=scaling,
                         rank_page_table=page_table,
                         rank_page_start=page_start,
                         ranks=ranks,
                         max_rank=max_rank)
        t.assert_close(gt, output, atol=2e-3, rtol=1e-5)
//  [python]
def test_mbgmv(self, input, paged_lora_a, paged_lora_b, out_head_size,
                   adapter_ids, scaling, page_table, ranks, page_start, gt):
        max_rank = page_table.size(-1)

        xa = mbgmv_a(input,
                     paged_lora_a,
                     adapter_ids=adapter_ids,
                     rank_page_table=page_table,
                     rank_page_start=page_start,
                     ranks=ranks,
                     max_rank=max_rank)

        output = mbgmv_b(xa,
                         paged_lora_b[..., :out_head_size],
                         adapter_ids=adapter_ids,
                         scaling=scaling,
                         rank_page_table=page_table,
                         rank_page_start=page_start,
                         ranks=ranks,
                         max_rank=max_rank)
        torch.testing.assert_close(gt, output, atol=2e-3, rtol=1e-5)
//  [python]
def assert_allclose(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    equal_nan: bool = True,
    msg: str = "",
) -> None:
    """
    .. warning::

       :func:`torch.testing.assert_allclose` is deprecated since ``1.12`` and will be removed in a future release.
       Please use :func:`torch.testing.assert_close` instead. You can find detailed upgrade instructions
       `here <https://github.com/pytorch/pytorch/issues/61844>`_.
    """
    warnings.warn(
        "`torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. "
        "Please use `torch.testing.assert_close()` instead. "
        "You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.",
        FutureWarning,
        stacklevel=2,
    )

    if not isinstance(actual, torch.Tensor):
        actual = torch.tensor(actual)
    if not isinstance(expected, torch.Tensor):
        expected = torch.tensor(expected, dtype=actual.dtype)

    if rtol is None and atol is None:
        rtol, atol = default_tolerances(
            actual,
            expected,
            dtype_precisions={
                torch.float16: (1e-3, 1e-3),
                torch.float32: (1e-4, 1e-5),
                torch.float64: (1e-5, 1e-8),
            },
        )

    t.assert_allclose(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
    )
//  [python]
def assert_allclose(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    equal_nan: bool = True,
    msg: str = "",
) -> None:
    """
    .. warning::

       :func:`torch.testing.assert_allclose` is deprecated since ``1.12`` and will be removed in a future release.
       Please use :func:`torch.testing.assert_close` instead. You can find detailed upgrade instructions
       `here <https://github.com/pytorch/pytorch/issues/61844>`_.
    """
    warnings.warn(
        "`torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. "
        "Please use `torch.testing.assert_close()` instead. "
        "You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.",
        FutureWarning,
        stacklevel=2,
    )

    if not isinstance(actual, torch.Tensor):
        actual = torch.tensor(actual)
    if not isinstance(expected, torch.Tensor):
        expected = torch.tensor(expected, dtype=actual.dtype)

    if rtol is None and atol is None:
        rtol, atol = default_tolerances(
            actual,
            expected,
            dtype_precisions={
                torch.float16: (1e-3, 1e-3),
                torch.float32: (1e-4, 1e-5),
                torch.float64: (1e-5, 1e-8),
            },
        )

    torch.testing.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
        check_device=True,
        check_dtype=False,
        check_stride=False,
        msg=msg or None,
    )
//  [python]
def test_elementwise_model(self):
        torch._dynamo.reset()
        """Test DORT with a pure function."""

        def run_elementwise_model():
            # A function to test DORT.
            def elementwise_model(tensor_x: torch.Tensor):
                tensor_w = tensor_x.sigmoid()
                tensor_y = tensor_w * tensor_w + 1.5
                tensor_z = tensor_y + tensor_x
                tensor_p = tensor_z * tensor_x
                tensor_q = tensor_p.sigmoid()
                return tensor_q

            optimized_elementwise_model = torch.compile(elementwise_model, backend="onnxrt", dynamic=True)

            def run(fun, list_x):
                tensor_x = torch.tensor(list_x, dtype=torch.float32).requires_grad_()
                tensor_y = fun(tensor_x)
                tensor_y.sum().backward()
                return tensor_x, tensor_y, tensor_x.grad

            # Baseline.
            tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, [-1.0, 2.0])
            # ORT result.
            tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, [-1.0, 2.0])

            torch.testing.assert_allclose(tensor_x, tensor_x_new)
            torch.testing.assert_allclose(tensor_y, tensor_y_new)
            torch.testing.assert_allclose(tensor_x_grad, tensor_x_grad_new)

        # Run 5 times because ORT runs have side effects and we want to make sure
        # the code is correct with them.
        for _ in range(5):
            run_elementwise_model()
//  [python]
def test_elementwise_model(self):
        torch._dynamo.reset()
        """Test DORT with a pure function."""

        def run_elementwise_model():
            # A function to test DORT.
            def elementwise_model(tensor_x: torch.Tensor):
                tensor_w = tensor_x.sigmoid()
                tensor_y = tensor_w * tensor_w + 1.5
                tensor_z = tensor_y + tensor_x
                tensor_p = tensor_z * tensor_x
                tensor_q = tensor_p.sigmoid()
                return tensor_q

            optimized_elementwise_model = torch.compile(elementwise_model, backend="onnxrt", dynamic=True)

            def run(fun, list_x):
                tensor_x = torch.tensor(list_x, dtype=torch.float32).requires_grad_()
                tensor_y = fun(tensor_x)
                tensor_y.sum().backward()
                return tensor_x, tensor_y, tensor_x.grad

            # Baseline.
            tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, [-1.0, 2.0])
            # ORT result.
            tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, [-1.0, 2.0])

            torch.testing.assert_close(tensor_x, tensor_x_new)
            torch.testing.assert_close(tensor_y, tensor_y_new)
            torch.testing.assert_close(tensor_x_grad, tensor_x_grad_new)

        # Run 5 times because ORT runs have side effects and we want to make sure
        # the code is correct with them.
        for _ in range(5):
            run_elementwise_model()
//  [python]
def test_aten_full(self):
        torch._dynamo.reset()

        def run_no_input_model():
            # A function to test.
            def no_input_model():
                return torch.ops.aten.full([2, 3], 1.5)

            @torch._dynamo.optimize("onnxrt")
            def optimized_no_input_model():
                return no_input_model()

            def run(fun):
                tensor_x = fun()
                return tensor_x

            # Baseline.
            tensor_x = run(no_input_model)
            # ORT result.
            tensor_x_new = run(optimized_no_input_model)

            torch.testing.assert_allclose(tensor_x, tensor_x_new)

        for _ in range(5):
            run_no_input_model()
//  [python]
def test_aten_full(self):
        torch._dynamo.reset()

        def run_no_input_model():
            # A function to test.
            def no_input_model():
                return torch.ops.aten.full([2, 3], 1.5)

            @torch._dynamo.optimize("onnxrt")
            def optimized_no_input_model():
                return no_input_model()

            def run(fun):
                tensor_x = fun()
                return tensor_x

            # Baseline.
            tensor_x = run(no_input_model)
            # ORT result.
            tensor_x_new = run(optimized_no_input_model)

            torch.testing.assert_close(tensor_x, tensor_x_new)

        for _ in range(5):
            run_no_input_model()
//  [python]
def test_aten_full_with_device(self):
        torch._dynamo.reset()

        def run_no_input_model():
            # A function to test.
            def no_input_model():
                return torch.ops.aten.full([2, 3], 1.5, device="cpu")

            optimized_no_input_model = torch.compile(no_input_model, backend="onnxrt")

            def run(fun):
                tensor_x = fun()
                return tensor_x

            # Baseline.
            tensor_x = run(no_input_model)
            # ORT result.
            tensor_x_new = run(optimized_no_input_model)

            torch.testing.assert_allclose(tensor_x, tensor_x_new)

        for _ in range(5):
            run_no_input_model()
//  [python]
def test_aten_full_with_device(self):
        torch._dynamo.reset()

        def run_no_input_model():
            # A function to test.
            def no_input_model():
                return torch.ops.aten.full([2, 3], 1.5, device="cpu")

            optimized_no_input_model = torch.compile(no_input_model, backend="onnxrt")

            def run(fun):
                tensor_x = fun()
                return tensor_x

            # Baseline.
            tensor_x = run(no_input_model)
            # ORT result.
            tensor_x_new = run(optimized_no_input_model)

            torch.testing.assert_close(tensor_x, tensor_x_new)

        for _ in range(5):
            run_no_input_model()
//  [python]
def test_upper_tail():
    t.assert_close(
        upper_tail(
            Normal(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.6745]),
    )

    t.assert_close(
        upper_tail(
            Uniform(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.75]),
    )
//  [python]
def test_upper_tail():
    torch.testing.assert_close(
        upper_tail(
            Normal(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.6745]),
    )

    torch.testing.assert_close(
        upper_tail(
            Uniform(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.75]),
    )
//  [python]
def test_elementwise_model_for_inference(self):
        torch._dynamo.reset()

        # A function to test DORT for inference (i.e., the compiled function
        # doesn't have backward pass).
        def elementwise_model(tensor_x: torch.Tensor):
            tensor_w = tensor_x.relu()
            tensor_y = tensor_w * tensor_w + 1.5
            tensor_z = tensor_y + tensor_x
            tensor_p = tensor_z * tensor_x
            tensor_q = tensor_p.relu()
            return tensor_q

        # TODO: Set use_aot_autograd=False. In order to decompose torch
        # function calls to aten ops, we need to set
        # user_aot_autograd=True because there is no decomposition in DORT
        # anymore. A long-term fix will be brining # decomposition pass back
        # into DORT.
        local_backend = make_local_backend(dynamic=True, use_aot_autograd=True)
        optimized_elementwise_model = torch.compile(elementwise_model, backend=local_backend, dynamic=True)

        def run(fun, list_x):
            tensor_x = torch.tensor(list_x, dtype=torch.float32).requires_grad_()
            tensor_y = fun(tensor_x)
            return tensor_y

        # Baseline.
        tensor_y = run(elementwise_model, [-1.0, 2.0])
        # ORT result.
        tensor_y_new = run(optimized_elementwise_model, [-1.0, 2.0])

        torch.testing.assert_allclose(tensor_y, tensor_y_new)
//  [python]
def test_elementwise_model_for_inference(self):
        torch._dynamo.reset()

        # A function to test DORT for inference (i.e., the compiled function
        # doesn't have backward pass).
        def elementwise_model(tensor_x: torch.Tensor):
            tensor_w = tensor_x.relu()
            tensor_y = tensor_w * tensor_w + 1.5
            tensor_z = tensor_y + tensor_x
            tensor_p = tensor_z * tensor_x
            tensor_q = tensor_p.relu()
            return tensor_q

        # TODO: Set use_aot_autograd=False. In order to decompose torch
        # function calls to aten ops, we need to set
        # user_aot_autograd=True because there is no decomposition in DORT
        # anymore. A long-term fix will be brining # decomposition pass back
        # into DORT.
        local_backend = make_local_backend(dynamic=True, use_aot_autograd=True)
        optimized_elementwise_model = torch.compile(elementwise_model, backend=local_backend, dynamic=True)

        def run(fun, list_x):
            tensor_x = torch.tensor(list_x, dtype=torch.float32).requires_grad_()
            tensor_y = fun(tensor_x)
            return tensor_y

        # Baseline.
        tensor_y = run(elementwise_model, [-1.0, 2.0])
        # ORT result.
        tensor_y_new = run(optimized_elementwise_model, [-1.0, 2.0])

        torch.testing.assert_close(tensor_y, tensor_y_new)
//  [python]
def test_to_copy(self):
        torch._dynamo.reset()
        """Test DORT with aten::_to_copy."""

        def run_to_copy():
            # A function to test.
            def copy_copy_copy(tensor_x: torch.Tensor):
                tensor_x1 = torch.ops.aten._to_copy(tensor_x, dtype=torch.int64)
                tensor_x2 = torch.ops.aten._to_copy(tensor_x, dtype=torch.int64, device=tensor_x.device)
                tensor_x3 = torch.ops.aten._to_copy(
                    tensor_x, dtype=torch.int64, device=tensor_x.device, layout=torch.strided
                )
                return tensor_x1, tensor_x2, tensor_x3

            optimized_copy_copy_copy = torch.compile(copy_copy_copy, backend="onnxrt")

            def run(fun, list_x):
                tensor_x = torch.tensor(list_x, dtype=torch.float32)
                tensor_x1, tensor_x2, tensor_x3 = fun(tensor_x)
                return tensor_x1, tensor_x2, tensor_x3

            # Baseline.
            tensor_x, tensor_y, tensor_z = run(copy_copy_copy, [-1.0, 2.0])
            # ORT result.
            tensor_x_new, tensor_y_new, tensor_z_new = run(optimized_copy_copy_copy, [-1.0, 2.0])

            torch.testing.assert_allclose(tensor_x, tensor_x_new)
            torch.testing.assert_allclose(tensor_y, tensor_y_new)
            torch.testing.assert_allclose(tensor_z, tensor_z_new)

        run_to_copy()
//  [python]
def test_to_copy(self):
        torch._dynamo.reset()
        """Test DORT with aten::_to_copy."""

        def run_to_copy():
            # A function to test.
            def copy_copy_copy(tensor_x: torch.Tensor):
                tensor_x1 = torch.ops.aten._to_copy(tensor_x, dtype=torch.int64)
                tensor_x2 = torch.ops.aten._to_copy(tensor_x, dtype=torch.int64, device=tensor_x.device)
                tensor_x3 = torch.ops.aten._to_copy(
                    tensor_x, dtype=torch.int64, device=tensor_x.device, layout=torch.strided
                )
                return tensor_x1, tensor_x2, tensor_x3

            optimized_copy_copy_copy = torch.compile(copy_copy_copy, backend="onnxrt")

            def run(fun, list_x):
                tensor_x = torch.tensor(list_x, dtype=torch.float32)
                tensor_x1, tensor_x2, tensor_x3 = fun(tensor_x)
                return tensor_x1, tensor_x2, tensor_x3

            # Baseline.
            tensor_x, tensor_y, tensor_z = run(copy_copy_copy, [-1.0, 2.0])
            # ORT result.
            tensor_x_new, tensor_y_new, tensor_z_new = run(optimized_copy_copy_copy, [-1.0, 2.0])

            torch.testing.assert_close(tensor_x, tensor_x_new)
            torch.testing.assert_close(tensor_y, tensor_y_new)
            torch.testing.assert_close(tensor_z, tensor_z_new)

        run_to_copy()
//  [python]
def test_dynamo_shape_model(self):
        torch._dynamo.reset()
        """Test DORT with a pure function."""

        def run_elementwise_model():
            # A function to test DORT.
            def elementwise_model(tensor_x: torch.Tensor):
                tensor_y = tensor_x.sigmoid()
                tensor_z = tensor_y + tensor_x
                tensor_p = tensor_z * tensor_x
                tensor_q = tensor_p.sigmoid()
                return tensor_q

            # This function should only generate one graph and execute
            # it for all inputs.
            # With dynamic_shape=True, Dynamo sends FX graphs with dynamic
            # shapes (e.g., batch size is a symbol "batch" instead of a fixed
            # number) to OrtBackend.compile(...).
            optimized_elementwise_model = torch.compile(elementwise_model, backend="onnxrt", dynamic=True)

            def run(fun, seed: torch.Tensor):
                tensor_x = seed.detach().clone().requires_grad_()
                tensor_y = fun(tensor_x)
                tensor_y.sum().backward()
                return tensor_x, tensor_y, tensor_x.grad

            # Dimension changed.
            for shape in [(2, 3), (3, 4)]:
                seed = torch.rand(shape)
                # Baseline.
                tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, seed)
                # ORT result.
                tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, seed)

                torch.testing.assert_allclose(tensor_x, tensor_x_new)
                torch.testing.assert_allclose(tensor_y, tensor_y_new)
                torch.testing.assert_allclose(tensor_x_grad, tensor_x_grad_new)

            # Rank changed.
            for shape in [(1,), (2,), (2, 3), (2, 3, 4)]:
                seed = torch.rand(shape)
                # Baseline.
                tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, seed)
                # ORT result.
                tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, seed)

                torch.testing.assert_allclose(tensor_x, tensor_x_new)
                torch.testing.assert_allclose(tensor_y, tensor_y_new)
                torch.testing.assert_allclose(tensor_x_grad, tensor_x_grad_new)

        run_elementwise_model()
//  [python]
def test_dynamo_shape_model(self):
        torch._dynamo.reset()
        """Test DORT with a pure function."""

        def run_elementwise_model():
            # A function to test DORT.
            def elementwise_model(tensor_x: torch.Tensor):
                tensor_y = tensor_x.sigmoid()
                tensor_z = tensor_y + tensor_x
                tensor_p = tensor_z * tensor_x
                tensor_q = tensor_p.sigmoid()
                return tensor_q

            # This function should only generate one graph and execute
            # it for all inputs.
            # With dynamic_shape=True, Dynamo sends FX graphs with dynamic
            # shapes (e.g., batch size is a symbol "batch" instead of a fixed
            # number) to OrtBackend.compile(...).
            optimized_elementwise_model = torch.compile(elementwise_model, backend="onnxrt", dynamic=True)

            def run(fun, seed: torch.Tensor):
                tensor_x = seed.detach().clone().requires_grad_()
                tensor_y = fun(tensor_x)
                tensor_y.sum().backward()
                return tensor_x, tensor_y, tensor_x.grad

            # Dimension changed.
            for shape in [(2, 3), (3, 4)]:
                seed = torch.rand(shape)
                # Baseline.
                tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, seed)
                # ORT result.
                tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, seed)

                torch.testing.assert_close(tensor_x, tensor_x_new)
                torch.testing.assert_close(tensor_y, tensor_y_new)
                torch.testing.assert_close(tensor_x_grad, tensor_x_grad_new)

            # Rank changed.
            for shape in [(1,), (2,), (2, 3), (2, 3, 4)]:
                seed = torch.rand(shape)
                # Baseline.
                tensor_x, tensor_y, tensor_x_grad = run(elementwise_model, seed)
                # ORT result.
                tensor_x_new, tensor_y_new, tensor_x_grad_new = run(optimized_elementwise_model, seed)

                torch.testing.assert_close(tensor_x, tensor_x_new)
                torch.testing.assert_close(tensor_y, tensor_y_new)
                torch.testing.assert_close(tensor_x_grad, tensor_x_grad_new)

        run_elementwise_model()
//  [python]
def test_pad(self):
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        img = torch.ones(3, height, width)
        padding = random.randint(1, 20)
        fill = random.randint(1, 50)
        result = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Pad(padding, fill=fill),
            transforms.ToTensor(),
        ])(img)
        self.assertEqual(result.size(1), height + 2 * padding)
        self.assertEqual(result.size(2), width + 2 * padding)
        # check that all elements in the padded region correspond
        # to the pad value
        fill_v = fill / 255
        eps = 1e-5
        h_padded = result[:, :padding, :]
        w_padded = result[:, :, :padding]
        self.assertTrue(torch.allclose(
            h_padded, torch.full_like(h_padded, fill_value=fill_v), rtol=0.0, atol=eps
        ))
        self.assertTrue(torch.allclose(
            w_padded, torch.full_like(w_padded, fill_value=fill_v), rtol=0.0, atol=eps
        ))
        self.assertRaises(ValueError, transforms.Pad(padding, fill=(1, 2)),
                          transforms.ToPILImage()(img))
//  [python]
def test_pad(self):
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        img = torch.ones(3, height, width)
        padding = random.randint(1, 20)
        fill = random.randint(1, 50)
        result = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Pad(padding, fill=fill),
            transforms.ToTensor(),
        ])(img)
        self.assertEqual(result.size(1), height + 2 * padding)
        self.assertEqual(result.size(2), width + 2 * padding)
        # check that all elements in the padded region correspond
        # to the pad value
        fill_v = fill / 255
        eps = 1e-5
        h_padded = result[:, :padding, :]
        w_padded = result[:, :, :padding]
        torch.testing.assert_close(
            h_padded, torch.full_like(h_padded, fill_value=fill_v), check_stride=False, rtol=0.0, atol=eps
        )
        torch.testing.assert_close(
            w_padded, torch.full_like(w_padded, fill_value=fill_v), check_stride=False, rtol=0.0, atol=eps
        )
        self.assertRaises(ValueError, transforms.Pad(padding, fill=(1, 2)),
                          transforms.ToPILImage()(img))
//  [python]
def test_accimage_to_tensor(self):
        trans = transforms.ToTensor()

        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.ToTensor()

        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_convert_image_dtype_float_to_float(self):
        for input_dtype, output_dtypes in cycle_over(float_dtypes()):
            input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
            for output_dtype in output_dtypes:
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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)

                    self.assertEqual(output_image_script.dtype, output_dtype)
                    self.assertEqual(output_image.dtype, output_dtype)
                    self.assertTrue(torch.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

                    self.assertAlmostEqual(actual_min, desired_min)
                    self.assertAlmostEqual(actual_max, desired_max)
//  [python]
def test_convert_image_dtype_float_to_float(self):
        for input_dtype, output_dtypes in cycle_over(float_dtypes()):
            input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
            for output_dtype in output_dtypes:
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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

                    self.assertAlmostEqual(actual_min, desired_min)
                    self.assertAlmostEqual(actual_max, desired_max)
//  [python]
def test_convert_image_dtype_float_to_int(self):
        for input_dtype in float_dtypes():
            input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
            for output_dtype in int_dtypes():
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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 self.assertRaises(RuntimeError):
                            transform(input_image)
                    else:
                        output_image = transform(input_image)
                        output_image_script = transform_script(input_image, output_dtype)

                        self.assertEqual(output_image_script.tolist(), output_image.tolist())

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

                        self.assertEqual(actual_min, desired_min)
                        self.assertEqual(actual_max, desired_max)
//  [python]
def test_convert_image_dtype_float_to_int(self):
        for input_dtype in float_dtypes():
            input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)
            for output_dtype in int_dtypes():
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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 self.assertRaises(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

                        self.assertEqual(actual_min, desired_min)
                        self.assertEqual(actual_max, desired_max)
//  [python]
def test_convert_image_dtype_int_to_float(self):
        for input_dtype in int_dtypes():
            input_image = torch.tensor((0, torch.iinfo(input_dtype).max), dtype=input_dtype)
            for output_dtype in float_dtypes():
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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)

                    self.assertEqual(output_image_script.dtype, output_dtype)
                    self.assertEqual(output_image.dtype, 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

                    self.assertAlmostEqual(actual_min, desired_min)
                    self.assertGreaterEqual(actual_min, desired_min)
                    self.assertAlmostEqual(actual_max, desired_max)
                    self.assertLessEqual(actual_max, desired_max)
//  [python]
def test_convert_image_dtype_int_to_float(self):
        for input_dtype in int_dtypes():
            input_image = torch.tensor((0, torch.iinfo(input_dtype).max), dtype=input_dtype)
            for output_dtype in float_dtypes():
                with self.subTest(input_dtype=input_dtype, output_dtype=output_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

                    self.assertAlmostEqual(actual_min, desired_min)
                    self.assertGreaterEqual(actual_min, desired_min)
                    self.assertAlmostEqual(actual_max, desired_max)
                    self.assertLessEqual(actual_max, desired_max)