jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_delete_unused_submodules_leaf(self):
        class SubModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(10, 10)
                self.relu = torch.nn.ReLU()

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

        class Model(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.submod = SubModule()

            def forward(self, x):
                x = self.submod(x)
                return x

        model = Model()

        class MyCustomTracer(torch.fx.Tracer):
            def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> bool:
                return module_qualified_name == "submod"

        inputs = torch.randn(1, 10)
        traced_graph = MyCustomTracer().trace(model)
        gm2 = torch.fx.GraphModule(model, traced_graph)
        gm2.delete_all_unused_submodules()
        torch.testing.assert_close(gm2(inputs), model(inputs))
//  [python]
def test_load_mobile_module_with_debug_info(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, x):
                return x + 5

        input = torch.tensor([3])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter(
                _save_mobile_debug_info=True
            )
        )
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        t.assert_close(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        t.assert_close(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        t.assert_close(
            script_module_result, mobile_module_run_method_result
        )
//  [python]
def test_load_mobile_module_with_debug_info(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, x):
                return x + 5

        input = torch.tensor([3])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter(
                _save_mobile_debug_info=True
            )
        )
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        torch.testing.assert_close(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_close(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_close(
            script_module_result, mobile_module_run_method_result
        )
//  [python]
def test_find_and_run_method(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, arg):
                return arg

        input = (torch.tensor([1]),)

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(*input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertFalse(has_bundled_inputs)

        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            script_module, [input], []
        )

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertTrue(has_bundled_inputs)

        bundled_inputs = mobile_module.run_method("get_all_bundled_inputs")
        mobile_module_result = mobile_module.forward(*bundled_inputs[0])
        self.assertEqual(script_module_result, mobile_module_result)
//  [python]
def test_find_and_run_method(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, arg):
                return arg

        input = (torch.tensor([1]),)

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(*input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertFalse(has_bundled_inputs)

        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            script_module, [input], []
        )

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertTrue(has_bundled_inputs)

        bundled_inputs = mobile_module.run_method("get_all_bundled_inputs")
        mobile_module_result = mobile_module.forward(*bundled_inputs[0])
        torch.testing.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_benchmark_geglu(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = GEGLU(in_features, out_features,
                  bias=bias).cuda().to(dtype=dtype).eval()
        x = torch.randn(N, in_features).cuda().to(dtype=dtype)
        out = m(x)
        out_opt = m(x, enable_opt=True)

        torch.testing.assert_allclose(out_opt, out, rtol=2e-2, atol=2e-2)

        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            out = m(x)
        torch.cuda.synchronize()
        logger.info(f'cost={time.time() - start}')

        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            out_opt = m(x, enable_opt=True)
        torch.cuda.synchronize()
        logger.info(f'cost={time.time() - start}')
//  [python]
def test_benchmark_geglu(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = GEGLU(in_features, out_features,
                  bias=bias).cuda().to(dtype=dtype).eval()
        x = torch.randn(N, in_features).cuda().to(dtype=dtype)
        out = m(x)
        out_opt = m(x, enable_opt=True)

        torch.testing.assert_close(out_opt, out, rtol=2e-2, atol=2e-2)

        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            out = m(x)
        torch.cuda.synchronize()
        logger.info(f'cost={time.time() - start}')

        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            out_opt = m(x, enable_opt=True)
        torch.cuda.synchronize()
        logger.info(f'cost={time.time() - start}')
//  [python]
def check_outputs(new_output, original_output):
        assert len(new_output) == 1
        torch.testing.assert_allclose(new_output[0], torch.tensor([5.0]))
//  [python]
def check_outputs(new_output, original_output):
        assert len(new_output) == 1
        torch.testing.assert_close(new_output[0], torch.tensor([5.0]))
//  [python]
def check_outputs(new_output, original_output):
        first, second = original_output
        t.assert_close(first["add"], torch.tensor([4]))
        t.assert_close(first["sub"], torch.tensor([-2]))
        t.assert_close(second["add"], torch.tensor([6]))
        t.assert_close(second["sub"], torch.tensor([-2]))
//  [python]
def check_outputs(new_output, original_output):
        first, second = original_output
        torch.testing.assert_close(first["add"], torch.tensor([4]))
        torch.testing.assert_close(first["sub"], torch.tensor([-2]))
        torch.testing.assert_close(second["add"], torch.tensor([6]))
        torch.testing.assert_close(second["sub"], torch.tensor([-2]))
//  [python]
def test_method_calls_with_optional_arg(self):
        class A(torch.nn.Module):
            def __init__(self):
                super().__init__()

            # opt arg in script-to-script invocation
            def forward(self, x, two: int = 2):
                return x + two

        class B(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.A0 = A()

            # opt arg in Python-to-script invocation
            def forward(self, x, one: int = 1):
                return self.A0(x) + one

        script_module = torch.jit.script(B())
        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        mobile_module = _load_for_lite_interpreter(buffer)

        input = torch.tensor([5])
        script_module_forward_result = script_module.forward(input)
        mobile_module_forward_result = mobile_module.forward(input)
        self.assertEqual(script_module_forward_result, mobile_module_forward_result)

        # change ref only
        script_module_forward_result = script_module.forward(input, 2)
        self.assertFalse(
            (script_module_forward_result == mobile_module_forward_result).all().item()
        )

        # now both match again
        mobile_module_forward_result = mobile_module.forward(input, 2)
        self.assertEqual(script_module_forward_result, mobile_module_forward_result)
//  [python]
def test_method_calls_with_optional_arg(self):
        class A(torch.nn.Module):
            def __init__(self):
                super().__init__()

            # opt arg in script-to-script invocation
            def forward(self, x, two: int = 2):
                return x + two

        class B(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.A0 = A()

            # opt arg in Python-to-script invocation
            def forward(self, x, one: int = 1):
                return self.A0(x) + one

        script_module = torch.jit.script(B())
        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        mobile_module = _load_for_lite_interpreter(buffer)

        input = torch.tensor([5])
        script_module_forward_result = script_module.forward(input)
        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_close(
            script_module_forward_result, mobile_module_forward_result
        )

        # change ref only
        script_module_forward_result = script_module.forward(input, 2)
        self.assertFalse(
            (script_module_forward_result == mobile_module_forward_result).all().item()
        )

        # now both match again
        mobile_module_forward_result = mobile_module.forward(input, 2)
        torch.testing.assert_close(
            script_module_forward_result, mobile_module_forward_result
        )
//  [python]
def test_broadcast_tensors(self, device):
        def fn(x):
            nz = x.nonzero()
            a = torch.zeros([nz.size(0), 512])
            b = torch.ones([nz.size(0), 1])
            return a * b

        x = torch.randn(32, 4, device=device)
        actual = torch.compile(fn, fullgraph=True)(x)
        expected = fn(x)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_broadcast_tensors(self, device):
        def fn(x):
            nz = x.nonzero()
            a = torch.zeros([nz.size(0), 512])
            b = torch.ones([nz.size(0), 1])
            return a * b

        x = torch.randn(32, 4, device=device)
        actual = torch.compile(fn, fullgraph=True)(x)
        expected = fn(x)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_expand(self, device):
        def fn(x, y):
            nz = torch.nonzero(x)
            # unbacked symint in nz.size
            x_exp = nz.expand([-1, 128])
            # unbacked symint in target sizes
            y_exp = y.expand([-1, nz.size(0)])
            return x_exp, y_exp

        example_inputs = (
            torch.randn((32), device=device),
            torch.randn((32, 1), device=device),
        )

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)

        self.assertEqual(actual, expected)
//  [python]
def test_expand(self, device):
        def fn(x, y):
            nz = torch.nonzero(x)
            # unbacked symint in nz.size
            x_exp = nz.expand([-1, 128])
            # unbacked symint in target sizes
            y_exp = y.expand([-1, nz.size(0)])
            return x_exp, y_exp

        example_inputs = (
            torch.randn((32), device=device),
            torch.randn((32, 1), device=device),
        )

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_autotuning(self, device):
        def fn(x, y):
            nz = torch.nonzero(x)
            # unbacked symint in the GEMM input shape
            a = x.new_ones([nz.size(0), y.size(0)])
            return a @ y

        example_inputs = (
            torch.randn((64), device=device),
            torch.randn((32, 16), device=device),
        )

        with inductor_config.patch(
            {
                "max_autotune_gemm": True,
            }
        ):
            actual = torch.compile(fn, fullgraph=True)(*example_inputs)
            expected = fn(*example_inputs)

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_autotuning(self, device):
        def fn(x, y):
            nz = torch.nonzero(x)
            # unbacked symint in the GEMM input shape
            a = x.new_ones([nz.size(0), y.size(0)])
            return a @ y

        example_inputs = (
            torch.randn((64), device=device),
            torch.randn((32, 16), device=device),
        )

        with inductor_config.patch(
            {
                "max_autotune_gemm": True,
            }
        ):
            actual = torch.compile(fn, fullgraph=True)(*example_inputs)
            expected = fn(*example_inputs)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_split_with_sizes(self, device):
        def fn(x, y):
            l = y.tolist()
            s = torch.split(x, l)
            d = l[0] + l[1] + l[2]
            return s[0].sum(), d

        example_inputs = (torch.randn((32), device=device), torch.tensor((7, 16, 9)))

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)

        self.assertEqual(actual, expected)
//  [python]
def test_split_with_sizes(self, device):
        def fn(x, y):
            l = y.tolist()
            s = torch.split(x, l)
            d = l[0] + l[1] + l[2]
            return s[0].sum(), d

        example_inputs = (torch.randn((32), device=device), torch.tensor((7, 16, 9)))

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_nonzero_in_inference_mode(self, device):
        def fn(x):
            return torch.nonzero(x)

        example_inputs = (torch.randint(0, 2, (128,), device=device),)

        with torch.inference_mode():
            actual = torch.compile(fn, fullgraph=True)(*example_inputs)
            expected = fn(*example_inputs)

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_nonzero_in_inference_mode(self, device):
        def fn(x):
            return torch.nonzero(x)

        example_inputs = (torch.randint(0, 2, (128,), device=device),)

        with torch.inference_mode():
            actual = torch.compile(fn, fullgraph=True)(*example_inputs)
            expected = fn(*example_inputs)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_view_of_slice(self, device):
        # Tests View.create(slice, size_with_unbacked_symint)
        def fn(x):
            nz = torch.nonzero(x)  # introduce unbacked symint
            squared = nz * nz  # avoid ReinterpretView when lowering Slice
            sliced = torch.ops.aten.slice.Tensor(squared, dim=1, start=-2, end=None)
            view = sliced.unsqueeze(dim=0)
            return view.squeeze(
                dim=0
            )  # make sure no unbacked symint in output's stride

        example_inputs = (torch.randn(1, 1, 1, 1, device=device),)
        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_view_of_slice(self, device):
        # Tests View.create(slice, size_with_unbacked_symint)
        def fn(x):
            nz = torch.nonzero(x)  # introduce unbacked symint
            squared = nz * nz  # avoid ReinterpretView when lowering Slice
            sliced = torch.ops.aten.slice.Tensor(squared, dim=1, start=-2, end=None)
            view = sliced.unsqueeze(dim=0)
            return view.squeeze(
                dim=0
            )  # make sure no unbacked symint in output's stride

        example_inputs = (torch.randn(1, 1, 1, 1, device=device),)
        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_triton_kernel_grid(self, device):
        if device == "cpu":
            raise unittest.SkipTest("Triton kernel requires GPU")

        from torch.testing._internal.triton_utils import add_kernel

        def fn(x):
            maxlen = max(x.item(), 512)
            a = torch.ones(maxlen, device=device)
            b = torch.ones(maxlen, device=device)
            out = torch.zeros_like(a)
            # unbacked symint in grid
            add_kernel[(1, 1, maxlen)](a, b, out, maxlen, 32)
            return out

        example_inputs = (torch.randint(high=1024, size=(1,), device=device),)
        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_triton_kernel_grid(self, device):
        if device == "cpu":
            raise unittest.SkipTest("Triton kernel requires GPU")

        from torch.testing._internal.triton_utils import add_kernel

        def fn(x):
            maxlen = max(x.item(), 512)
            a = torch.ones(maxlen, device=device)
            b = torch.ones(maxlen, device=device)
            out = torch.zeros_like(a)
            # unbacked symint in grid
            add_kernel[(1, 1, maxlen)](a, b, out, maxlen, 32)
            return out

        example_inputs = (torch.randint(high=1024, size=(1,), device=device),)
        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_word_information_lost(self) -> None:

        input = ["hello world", "welcome to the facebook"]
        target = ["hello metaverse", "welcome to meta"]
        t.assert_allclose(
            word_information_lost(input, target),
            torch.tensor(0.7, dtype=torch.float64),
        )

        input = ["this is the prediction", "there is an other sample"]
        target = ["this is the reference", "there is another one"]
        t.assert_allclose(
            word_information_lost(input, target),
            torch.tensor(0.6527777, dtype=torch.float64),
        )
//  [python]
def test_word_information_lost(self) -> None:

        input = ["hello world", "welcome to the facebook"]
        target = ["hello metaverse", "welcome to meta"]
        torch.testing.assert_close(
            word_information_lost(input, target),
            torch.tensor(0.7, dtype=torch.float64),
        )

        input = ["this is the prediction", "there is an other sample"]
        target = ["this is the reference", "there is another one"]
        torch.testing.assert_close(
            word_information_lost(input, target),
            torch.tensor(0.6527777, dtype=torch.float64),
        )
//  [python]
dist.init_process_group(rank=rank, world_size=world_size)

# Pipeline stage is our main pipeline runtime. It takes in the pipe object,
# the rank of this process, and the device.
stage = PipelineStage(pipe, rank, device)

# Input data
x = torch.randn(batch_size, d_hid, device=device)

# Run the pipeline with input `x`. Divide the batch into n micro-batches
# and run them in parallel on the pipeline
if rank == 0:
    stage(x)
elif rank == world_size - 1:
    output = stage()
else:
    stage()

if rank == world_size - 1:
    # Run the original code and get the output for comparison
    reference_output = model(x)
    # Compare numerics of pipeline and original model
    torch.testing.assert_allclose(output, reference_output)
    print(" Pipeline parallel model ran successfully! ".center(80, "*"))
//  [python]
dist.init_process_group(rank=rank, world_size=world_size)

# Pipeline stage is our main pipeline runtime. It takes in the pipe object,
# the rank of this process, and the device.
stage = PipelineStage(pipe, rank, device)

# Input data
x = torch.randn(batch_size, d_hid, device=device)

# Run the pipeline with input `x`. Divide the batch into n micro-batches
# and run them in parallel on the pipeline
if rank == 0:
    stage(x)
elif rank == world_size - 1:
    output = stage()
else:
    stage()

if rank == world_size - 1:
    # Run the original code and get the output for comparison
    reference_output = model(x)
    # Compare numerics of pipeline and original model
    torch.testing.assert_close(output, reference_output)
    print(" Pipeline parallel model ran successfully! ".center(80, "*"))
//  [python]
def test_equivalent_backed_unbacked(self, device):
        # Tests scenario when there are two equivalent backed & unbacked symints,
        # but when we look-up a size hint on the unbacked symint, we ignorantly
        # use the default fallback hint.

        def fn(x, w, a, b):
            # Make tensors where 1st dim is unbacked/backed.
            u0, s0 = a.item(), b.size(0)
            unbacked = x.expand(u0, *x.shape)
            backed = x.expand(s0, *x.shape)

            # The cat unifies u0 and s0 -- i.e. u0 == s0.
            cat = torch.cat([backed, unbacked, unbacked], dim=1)  # [s0, 30, 16]
            mat1 = torch.permute(cat, [0, 2, 1])  # [s0, 16, 30]
            mat2 = w.expand(u0, *w.shape)  # [u0, 30, 32]
            bmm = torch.ops.aten.bmm(mat1, mat2)
            return bmm

        example_inputs = (
            torch.randn((10, 16), dtype=torch.float32, device=device),
            torch.randn((30, 32), dtype=torch.float32, device=device),
            torch.tensor(7, device=device),
            backed := torch.randn((7,), device=device),
        )
        torch._dynamo.mark_dynamic(backed, 0)  # create backed symint

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_equivalent_backed_unbacked(self, device):
        # Tests scenario when there are two equivalent backed & unbacked symints,
        # but when we look-up a size hint on the unbacked symint, we ignorantly
        # use the default fallback hint.

        def fn(x, w, a, b):
            # Make tensors where 1st dim is unbacked/backed.
            u0, s0 = a.item(), b.size(0)
            unbacked = x.expand(u0, *x.shape)
            backed = x.expand(s0, *x.shape)

            # The cat unifies u0 and s0 -- i.e. u0 == s0.
            cat = torch.cat([backed, unbacked, unbacked], dim=1)  # [s0, 30, 16]
            mat1 = torch.permute(cat, [0, 2, 1])  # [s0, 16, 30]
            mat2 = w.expand(u0, *w.shape)  # [u0, 30, 32]
            bmm = torch.ops.aten.bmm(mat1, mat2)
            return bmm

        example_inputs = (
            torch.randn((10, 16), dtype=torch.float32, device=device),
            torch.randn((30, 32), dtype=torch.float32, device=device),
            torch.tensor(7, device=device),
            backed := torch.randn((7,), device=device),
        )
        torch._dynamo.mark_dynamic(backed, 0)  # create backed symint

        actual = torch.compile(fn, fullgraph=True)(*example_inputs)
        expected = fn(*example_inputs)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_get_prompt_logprobs(
    hf_runner,
    vllm_runner,
    model,
    dtype,
    example_prompts,
):
    max_tokens = 5
    hf_model = hf_runner(model, dtype=dtype)
    hf_logprobs = hf_model.generate_greedy_logprobs(
        example_prompts,
        max_tokens=max_tokens,
    )
    del hf_model

    vllm_model = vllm_runner(model, dtype=dtype)
    vllm_sampling_params = SamplingParams(max_tokens=max_tokens,
                                          logprobs=5,
                                          prompt_logprobs=5,
                                          temperature=0.0)
    vllm_results = vllm_model.model.generate(
        example_prompts, sampling_params=vllm_sampling_params)

    # Test whether logprobs are included in the results.
    for result in vllm_results:
        assert result.prompt_logprobs is not None
        assert result.outputs[0].logprobs is not None

    # Test whether prompt logprobs are consistent with HF
    for vllm_result, hf_logprob in zip(vllm_results, hf_logprobs):
        # Check prompt logprobs
        vllm_prompt_logprobs = vllm_result.prompt_logprobs[1:]
        for i, vllm_prompt_logprob_dict in enumerate(vllm_prompt_logprobs):
            for token_id, logprob in vllm_prompt_logprob_dict.items():
                t.assert_allclose(logprob,
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        vllm_sample_logprobs = vllm_result.outputs[0].logprobs
        for i, vllm_sample_logprob_dict in enumerate(vllm_sample_logprobs):
            for token_id, logprob in vllm_sample_logprob_dict.items():
                t.assert_allclose(logprob,
                                           hf_logprob[i][-1][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
//  [python]
def test_get_prompt_logprobs(
    hf_runner,
    vllm_runner,
    model,
    dtype,
    example_prompts,
):
    max_tokens = 5
    hf_model = hf_runner(model, dtype=dtype)
    hf_logprobs = hf_model.generate_greedy_logprobs(
        example_prompts,
        max_tokens=max_tokens,
    )
    del hf_model

    vllm_model = vllm_runner(model, dtype=dtype)
    vllm_sampling_params = SamplingParams(max_tokens=max_tokens,
                                          logprobs=5,
                                          prompt_logprobs=5,
                                          temperature=0.0)
    vllm_results = vllm_model.model.generate(
        example_prompts, sampling_params=vllm_sampling_params)

    # Test whether logprobs are included in the results.
    for result in vllm_results:
        assert result.prompt_logprobs is not None
        assert result.outputs[0].logprobs is not None

    # Test whether prompt logprobs are consistent with HF
    for vllm_result, hf_logprob in zip(vllm_results, hf_logprobs):
        # Check prompt logprobs
        vllm_prompt_logprobs = vllm_result.prompt_logprobs[1:]
        for i, vllm_prompt_logprob_dict in enumerate(vllm_prompt_logprobs):
            for token_id, logprob in vllm_prompt_logprob_dict.items():
                torch.testing.assert_close(logprob,
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        vllm_sample_logprobs = vllm_result.outputs[0].logprobs
        for i, vllm_sample_logprob_dict in enumerate(vllm_sample_logprobs):
            for token_id, logprob in vllm_sample_logprob_dict.items():
                torch.testing.assert_close(logprob,
                                           hf_logprob[i][-1][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
//  [python]
def test_lazy_trace():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))

        out = model(*args, **kwargs)

        model = lazy_trace(model)

        traced_out = model(*args, **kwargs)

        t.assert_close(out, traced_out)
//  [python]
def test_lazy_trace():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))

        out = model(*args, **kwargs)

        model = lazy_trace(model)

        traced_out = model(*args, **kwargs)

        torch.testing.assert_close(out, traced_out)
//  [python]
def test_epnp_6pts(self):
        # create some random test sample for a single camera
        pose = pp.SE3([ 0.0000, -8.0000,  0.0000,
                        0.0000, -0.3827,  0.0000,  0.9239]).to(torch.float64)
        f, img_size = 2, (9, 9)
        projection = torch.tensor([[f, 0, img_size[0] / 2],
                                   [0, f, img_size[1] / 2],
                                   [0, 0, 1              ]], dtype=torch.float64)
        # some random points in the view
        pts_c = torch.tensor([[2., 0., 2.],
                              [1., 0., 2.],
                              [0., 1., 1.],
                              [0., 0., 1.],
                              [1., 0., 1.],
                              [5., 5., 3.]], dtype=torch.float64)
        pixels = pp.homo2cart(pts_c @ projection.T)
        # transform the points to world coordinate
        # solve the PnP problem to find the camera pose
        pts_w = pose.Inv().Act(pts_c)
        # solve the PnP problem
        epnp = pp.module.EPnP(intrinsics=projection)
        solved_pose = epnp(pts_w, pixels)
        t.assert_close(solved_pose, pose, rtol=1e-4, atol=1e-4)
//  [python]
def test_epnp_6pts(self):
        # create some random test sample for a single camera
        pose = pp.SE3([ 0.0000, -8.0000,  0.0000,
                        0.0000, -0.3827,  0.0000,  0.9239]).to(torch.float64)
        f, img_size = 2, (9, 9)
        projection = torch.tensor([[f, 0, img_size[0] / 2],
                                   [0, f, img_size[1] / 2],
                                   [0, 0, 1              ]], dtype=torch.float64)
        # some random points in the view
        pts_c = torch.tensor([[2., 0., 2.],
                              [1., 0., 2.],
                              [0., 1., 1.],
                              [0., 0., 1.],
                              [1., 0., 1.],
                              [5., 5., 3.]], dtype=torch.float64)
        pixels = pp.homo2cart(pts_c @ projection.T)
        # transform the points to world coordinate
        # solve the PnP problem to find the camera pose
        pts_w = pose.Inv().Act(pts_c)
        # solve the PnP problem
        epnp = pp.module.EPnP(intrinsics=projection)
        solved_pose = epnp(pts_w, pixels)
        torch.testing.assert_close(solved_pose, pose, rtol=1e-4, atol=1e-4)
//  [python]
def test_lazy_trace_with_generator():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y,
                      alpha=0.5,
                      beta_gamma=(1, 0.5),
                      generator=torch.Generator())

        out, g = model(*args, **kwargs)

        model = lazy_trace(model, check_trace=False, strict=False)

        traced_out, traced_g = model(*args, **kwargs)
        # print(tuple(model._cached.values())[0].module.inlined_graph)

        torch.testing.assert_allclose(out, traced_out)
//  [python]
def test_lazy_trace_with_generator():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y,
                      alpha=0.5,
                      beta_gamma=(1, 0.5),
                      generator=torch.Generator())

        out, g = model(*args, **kwargs)

        model = lazy_trace(model, check_trace=False, strict=False)

        traced_out, traced_g = model(*args, **kwargs)
        # print(tuple(model._cached.values())[0].module.inlined_graph)

        torch.testing.assert_close(out, traced_out)
//  [python]
def test_ate_ite_cf_two_node(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    sem = create_lingauss_sem(two_variable_dict, coef_matrix, graph, log_scale=math.log(1e-8))
    intervention_values_a = TensorDict({"x2": torch.tensor([1.42, 0.42])}, batch_size=tuple())
    intervention_values_b = TensorDict({"x2": torch.tensor([0.42, 1.42])}, batch_size=tuple())
    average_treatment_effect = ate(sem, intervention_values_a, intervention_values_b)
    sample_size = 100
    factual_data = sem.sample(torch.Size([sample_size]))
    if graph[0, 1] > 0.0:
        expected_treatment_effect = torch.zeros_like(average_treatment_effect["x1"])
        expected_mean_a = factual_data["x1"]
    else:
        expected_mean_a = torch.einsum("i,ij->j", intervention_values_a["x2"], coef_matrix[1:, :1])
        expected_mean_b = torch.einsum("i,ij->j", intervention_values_b["x2"], coef_matrix[1:, :1])
        expected_treatment_effect = expected_mean_a - expected_mean_b
    torch.testing.assert_allclose(average_treatment_effect["x1"], expected_treatment_effect)

    individual_treatment_effect = ite(sem, factual_data, intervention_values_a, intervention_values_b)
    cf_effect = counterfactual(sem, factual_data, intervention_values_a)
    torch.testing.assert_allclose(individual_treatment_effect["x1"], expected_treatment_effect.expand((sample_size, 1)))
    torch.testing.assert_allclose(cf_effect["x1"], expected_mean_a.expand((sample_size, 1)))
//  [python]
def test_ate_ite_cf_two_node(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    sem = create_lingauss_sem(two_variable_dict, coef_matrix, graph, log_scale=math.log(1e-8))
    intervention_values_a = TensorDict({"x2": torch.tensor([1.42, 0.42])}, batch_size=tuple())
    intervention_values_b = TensorDict({"x2": torch.tensor([0.42, 1.42])}, batch_size=tuple())
    average_treatment_effect = ate(sem, intervention_values_a, intervention_values_b)
    sample_size = 100
    factual_data = sem.sample(torch.Size([sample_size]))
    if graph[0, 1] > 0.0:
        expected_treatment_effect = torch.zeros_like(average_treatment_effect["x1"])
        expected_mean_a = factual_data["x1"]
    else:
        expected_mean_a = torch.einsum("i,ij->j", intervention_values_a["x2"], coef_matrix[1:, :1])
        expected_mean_b = torch.einsum("i,ij->j", intervention_values_b["x2"], coef_matrix[1:, :1])
        expected_treatment_effect = expected_mean_a - expected_mean_b
    torch.testing.assert_close(average_treatment_effect["x1"], expected_treatment_effect)

    individual_treatment_effect = ite(sem, factual_data, intervention_values_a, intervention_values_b)
    cf_effect = counterfactual(sem, factual_data, intervention_values_a)
    torch.testing.assert_close(individual_treatment_effect["x1"], expected_treatment_effect.expand((sample_size, 1)))
    torch.testing.assert_close(cf_effect["x1"], expected_mean_a.expand((sample_size, 1)))
//  [python]
def test_ate_ite_cf_three_node(three_variable_dict):
    """x1->x2->x3"""
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[1, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph, log_scale=math.log(1e-8))
    intervention_values_a = TensorDict({"x2": torch.tensor([1.42, 0.42])}, batch_size=tuple())
    intervention_values_b = TensorDict({"x2": torch.tensor([0.42, 1.42])}, batch_size=tuple())
    average_treatment_effect = ate(sem, intervention_values_a, intervention_values_b)
    t.assert_close(average_treatment_effect["x1"], torch.zeros_like(average_treatment_effect["x1"]))
    expected_mean_a = torch.einsum("i,ij->j", intervention_values_a["x2"], coef_matrix[2:4, 4:])
    expected_mean_b = torch.einsum("i,ij->j", intervention_values_b["x2"], coef_matrix[2:4, 4:])
    expected_treatment_effect = expected_mean_a - expected_mean_b
    t.assert_close(average_treatment_effect["x3"], expected_treatment_effect)

    sample_size = 100
    factual_data = sem.sample(torch.Size([sample_size]))
    individual_treatment_effect = ite(sem, factual_data, intervention_values_a, intervention_values_b)
    cf_effect = counterfactual(sem, factual_data, intervention_values_a)
    t.assert_close(individual_treatment_effect["x1"], torch.zeros_like(individual_treatment_effect["x1"]))
    t.assert_close(individual_treatment_effect["x3"], expected_treatment_effect.expand((sample_size, 1)))
    t.assert_close(cf_effect["x1"], factual_data["x1"])
    t.assert_close(cf_effect["x3"], expected_mean_a.expand((sample_size, 1)))
//  [python]
def test_ate_ite_cf_three_node(three_variable_dict):
    """x1->x2->x3"""
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[1, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph, log_scale=math.log(1e-8))
    intervention_values_a = TensorDict({"x2": torch.tensor([1.42, 0.42])}, batch_size=tuple())
    intervention_values_b = TensorDict({"x2": torch.tensor([0.42, 1.42])}, batch_size=tuple())
    average_treatment_effect = ate(sem, intervention_values_a, intervention_values_b)
    torch.testing.assert_close(average_treatment_effect["x1"], torch.zeros_like(average_treatment_effect["x1"]))
    expected_mean_a = torch.einsum("i,ij->j", intervention_values_a["x2"], coef_matrix[2:4, 4:])
    expected_mean_b = torch.einsum("i,ij->j", intervention_values_b["x2"], coef_matrix[2:4, 4:])
    expected_treatment_effect = expected_mean_a - expected_mean_b
    torch.testing.assert_close(average_treatment_effect["x3"], expected_treatment_effect)

    sample_size = 100
    factual_data = sem.sample(torch.Size([sample_size]))
    individual_treatment_effect = ite(sem, factual_data, intervention_values_a, intervention_values_b)
    cf_effect = counterfactual(sem, factual_data, intervention_values_a)
    torch.testing.assert_close(individual_treatment_effect["x1"], torch.zeros_like(individual_treatment_effect["x1"]))
    torch.testing.assert_close(individual_treatment_effect["x3"], expected_treatment_effect.expand((sample_size, 1)))
    torch.testing.assert_close(cf_effect["x1"], factual_data["x1"])
    torch.testing.assert_close(cf_effect["x3"], expected_mean_a.expand((sample_size, 1)))
//  [python]
def test_trace_with_kwargs():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))
        traced_model, call_helper = trace_with_kwargs(model, args, kwargs)

        logging.info('Traced graph for eval:\n{}'.format(traced_model.graph))

        traced_model = call_helper(traced_model)
        out = model(*args, **kwargs)
        traced_out = traced_model(*args, **kwargs)

        t.assert_close(out, traced_out)

    model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)

    x = torch.ones(1, 1, 256, 256, requires_grad=True)
    y = torch.ones(1, 1, 254, 254, requires_grad=True)
    args = (x, )
    kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))
    traced_model, call_helper = trace_with_kwargs(model, args, kwargs)

    logging.info('Traced graph for training:\n{}'.format(traced_model.graph))

    out = model(*args, **kwargs)
    out.sum().backward()
    traced_out = call_helper(traced_model)(*args, **kwargs)
    traced_out.sum().backward()

    t.assert_close(traced_out, out)
    t.assert_close(traced_model.module.conv.weight.grad,
                                  model.conv.weight.grad)
    t.assert_close(traced_model.module.conv.bias.grad,
                                  model.conv.bias.grad)
//  [python]
def test_trace_with_kwargs():
    with torch.no_grad():
        model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)
        model.eval()

        x = torch.ones(1, 1, 256, 256)
        y = torch.ones(1, 1, 254, 254)
        args = (x, )
        kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))
        traced_model, call_helper = trace_with_kwargs(model, args, kwargs)

        logging.info('Traced graph for eval:\n{}'.format(traced_model.graph))

        traced_model = call_helper(traced_model)
        out = model(*args, **kwargs)
        traced_out = traced_model(*args, **kwargs)

        torch.testing.assert_close(out, traced_out)

    model = ConvBiasAddActivation(activation_cls=torch.nn.ReLU)

    x = torch.ones(1, 1, 256, 256, requires_grad=True)
    y = torch.ones(1, 1, 254, 254, requires_grad=True)
    args = (x, )
    kwargs = dict(y=y, alpha=0.5, beta_gamma=(1, 0.5))
    traced_model, call_helper = trace_with_kwargs(model, args, kwargs)

    logging.info('Traced graph for training:\n{}'.format(traced_model.graph))

    out = model(*args, **kwargs)
    out.sum().backward()
    traced_out = call_helper(traced_model)(*args, **kwargs)
    traced_out.sum().backward()

    torch.testing.assert_close(traced_out, out)
    torch.testing.assert_close(traced_model.module.conv.weight.grad,
                                  model.conv.weight.grad)
    torch.testing.assert_close(traced_model.module.conv.bias.grad,
                                  model.conv.bias.grad)
//  [python]
def check_dynamo(backend, device, err_msg):
    import torch

    if device == "cuda" and not torch.cuda.is_available():
        print(f"CUDA not available -- skipping CUDA check on {backend} backend\n")
        return

    try:
        import torch._dynamo as dynamo

        if device == "cuda":
            from torch.utils._triton import has_triton

            if not has_triton():
                print(
                    f"WARNING: CUDA available but triton cannot be used. "
                    f"Your GPU may not be supported. "
                    f"Skipping CUDA check on {backend} backend\n"
                )
                return

        dynamo.reset()

        @dynamo.optimize(backend, nopython=True)
        def fn(x):
            return x + x

        class Module(torch.nn.Module):
            def forward(self, x):
                return x + x

        mod = Module()
        opt_mod = dynamo.optimize(backend, nopython=True)(mod)

        for f in (fn, opt_mod):
            x = torch.randn(10, 10).to(device)
            x.requires_grad = True
            y = f(x)
            torch.testing.assert_allclose(y, x + x)
            z = y.sum()
            z.backward()
            torch.testing.assert_allclose(x.grad, 2 * torch.ones_like(x))
    except Exception:
        sys.stderr.write(traceback.format_exc() + "\n" + err_msg + "\n\n")
        sys.exit(1)
//  [python]
def check_dynamo(backend, device, err_msg):
    import torch

    if device == "cuda" and not torch.cuda.is_available():
        print(f"CUDA not available -- skipping CUDA check on {backend} backend\n")
        return

    try:
        import torch._dynamo as dynamo

        if device == "cuda":
            from torch.utils._triton import has_triton

            if not has_triton():
                print(
                    f"WARNING: CUDA available but triton cannot be used. "
                    f"Your GPU may not be supported. "
                    f"Skipping CUDA check on {backend} backend\n"
                )
                return

        dynamo.reset()

        @dynamo.optimize(backend, nopython=True)
        def fn(x):
            return x + x

        class Module(torch.nn.Module):
            def forward(self, x):
                return x + x

        mod = Module()
        opt_mod = dynamo.optimize(backend, nopython=True)(mod)

        for f in (fn, opt_mod):
            x = torch.randn(10, 10).to(device)
            x.requires_grad = True
            y = f(x)
            torch.testing.assert_close(y, x + x)
            z = y.sum()
            z.backward()
            torch.testing.assert_close(x.grad, 2 * torch.ones_like(x))
    except Exception:
        sys.stderr.write(traceback.format_exc() + "\n" + err_msg + "\n\n")
        sys.exit(1)
//  [python]
def testNpRelu(self):
        t.assert_allclose(
            np.array([[0., 0.7, 0.0, 0.3, 0.0], [0.1, 0.0, 0.5, 0.0, 0.9]]),
            self._npRelu(
                np.array([[-0.9, 0.7, -0.5, 0.3, -0.1], [0.1, -0.3, 0.5, -0.7,
                                                         0.9]])))
//  [python]
def testNpRelu(self):
        torch.testing.assert_close(
            np.array([[0., 0.7, 0.0, 0.3, 0.0], [0.1, 0.0, 0.5, 0.0, 0.9]]),
            self._npRelu(
                np.array([[-0.9, 0.7, -0.5, 0.3, -0.1], [0.1, -0.3, 0.5, -0.7,
                                                         0.9]])))
//  [python]
def testNpLeakyRelu(self):
        t.assert_close(
            np.array([[-0.09, 0.7, -0.05, 0.3, -0.01],
                      [0.1, -0.03, 0.5, -0.07, 0.9]]),
            self._npLeakyRelu(
                np.array([[-0.9, 0.7, -0.5, 0.3, -0.1], [0.1, -0.3, 0.5, -0.7,
                                                         0.9]]),
                negative_slope=0.1))
//  [python]
def testNpLeakyRelu(self):
        torch.testing.assert_close(
            np.array([[-0.09, 0.7, -0.05, 0.3, -0.01],
                      [0.1, -0.03, 0.5, -0.07, 0.9]]),
            self._npLeakyRelu(
                np.array([[-0.9, 0.7, -0.5, 0.3, -0.1], [0.1, -0.3, 0.5, -0.7,
                                                         0.9]]),
                negative_slope=0.1))
//  [python]
def test_flash_attn_llama_pattern():
    torch.cuda.empty_cache()
    torch.manual_seed(0)
    config = transformers.LlamaConfig(num_hidden_layers=2)
    config.hidden_size=512
    model = transformers.LlamaForCausalLM(config=config).to("cuda")

    flash_model = transformers.LlamaForCausalLM(config=config).to("cuda")
    flash_model.load_state_dict(model.state_dict())

    data = make_data(model, "cuda")
    flash_data = copy_nested_tensor(data)

    model.zero_grad(set_to_none=True)
    flash_model.zero_grad(set_to_none=True)
    flash_attn_gm = torch.compile(flash_model, backend=trival_compile_fx)
    with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16):
        input_data, label = data
        output = flash_model(input_data)
        golden_logits = output.logits
        golden_loss = F.cross_entropy(golden_logits.view(-1, model.config.vocab_size), label.view(-1))

        flash_input_data, flash_label = flash_data
        flash_output = flash_attn_gm(flash_input_data)
        flash_logits = flash_output["logits"]
        flash_loss = F.cross_entropy(flash_logits.view(-1, model.config.vocab_size), flash_label.view(-1))

    # Flash attention uses a different mask than the original llama, relax check
    torch.allclose(golden_loss, flash_loss, atol=3e-4, rtol=1e-6)
    torch.allclose(golden_logits, flash_logits, atol=3e-2, rtol=1e-6)