jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_fx2ait_args(self):
        class TestModule(torch.nn.Module):
            def forward(self, a, b, c, d):
                temp = a + b[0] + b[1] + c[0] + c[1] + d
                return temp

        mod = TestModule().half().cuda()

        a = torch.randn(5, 3).half().cuda()
        b = [torch.randn(5, 3).half().cuda(), torch.randn(5, 3).half().cuda()]
        c = [torch.randn(5, 3).half().cuda(), torch.randn(5, 3).half().cuda()]
        d = torch.randn(5, 3).half().cuda()
        ref_output = mod(a, b, c, d)

        traced = acc_tracer.trace(mod, [a, b, c, d])

        ait_dump_dir = tempfile.mkdtemp(prefix="test_fx2ait_", dir="/tmp")

        interp = AITInterpreter(traced, [a, b, c, d], ait_dump_dir, "test")
        interp_result = interp.run()
        ait_mod = AITModule(
            AIT_MODEL_CLASS(
                interp_result.engine.lib_path,
                interp_result.input_names,
                interp_result.output_names,
                torch.float16,
                torch.float16,
                1,  # num_runtimes
            ),
            interp_result,
        )
        ait_output = ait_mod(a, b, c, d)
        torch.testing.assert_close(ait_output, ref_output, atol=0.1, rtol=0.1)

        ait_output = ait_mod(a, b, c, d=d)
        torch.testing.assert_close(ait_output, ref_output, atol=0.1, rtol=0.1)

        ait_output = ait_mod(a, b, c=c, d=d)
        torch.testing.assert_close(ait_output, ref_output, atol=0.1, rtol=0.1)

        ait_output = ait_mod(a, b=b, c=c, d=d)
        torch.testing.assert_close(ait_output, ref_output, atol=0.1, rtol=0.1)
//  [python]
def test_remove_mlp_layer_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        output_vanish=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.remove_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size > 1:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_allclose(param, initial_net_dict[key]), evolvable_mlp
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_remove_mlp_layer_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        output_vanish=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.remove_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size > 1:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key]), evolvable_mlp
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_remove_mlp_layer(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        min_hidden_layers=1,
        max_hidden_layers=10,
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.remove_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size > 1:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            print(key, param.shape)
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.testing.assert_allclose(param, initial_net_dict[key]), evolvable_mlp
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_remove_mlp_layer(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        min_hidden_layers=1,
        max_hidden_layers=10,
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.remove_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size > 1:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            print(key, param.shape)
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.testing.assert_close(param, initial_net_dict[key]), evolvable_mlp
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def _test_split_tensor_block_recovery(
        self,
        original_tensor: Tensor,
        expected_split_tensors: List[Tensor],
        start_idx: int,
        end_idx: int,
    ) -> None:
        actual_split_tensors = FSDPDistributor._split_tensor_block_recovery(
            original_tensor.flatten()[start_idx:end_idx],
            original_tensor.size(),
            start_idx,
            end_idx,
        )

        self.assertNotEqual(len(actual_split_tensors), 0)
        self.assertEqual(len(actual_split_tensors), len(expected_split_tensors))
        for i in range(len(actual_split_tensors)):
            self.assertTrue(torch.allclose(actual_split_tensors[i], expected_split_tensors[i]))
//  [python]
def _test_split_tensor_block_recovery(
        self,
        original_tensor: Tensor,
        expected_split_tensors: List[Tensor],
        start_idx: int,
        end_idx: int,
    ) -> None:
        actual_split_tensors = FSDPDistributor._split_tensor_block_recovery(
            original_tensor.flatten()[start_idx:end_idx],
            original_tensor.size(),
            start_idx,
            end_idx,
        )

        self.assertNotEqual(len(actual_split_tensors), 0)
        torch.testing.assert_close(actual_split_tensors, expected_split_tensors)
//  [python]
def test_random_horizontal_flip_video(self, p):
        clip = torch.rand((3, 4, 112, 112), dtype=torch.float)
        hclip = clip.flip(-1)

        out = transforms.RandomHorizontalFlipVideo(p=p)(clip)
        if p == 0:
            torch.testing.assert_allclose(out, clip)
        elif p == 1:
            torch.testing.assert_allclose(out, hclip)

        transforms.RandomHorizontalFlipVideo().__repr__()
//  [python]
def test_random_horizontal_flip_video(self, p):
        clip = torch.rand((3, 4, 112, 112), dtype=torch.float)
        hclip = clip.flip(-1)

        out = transforms.RandomHorizontalFlipVideo(p=p)(clip)
        if p == 0:
            torch.testing.assert_close(out, clip)
        elif p == 1:
            torch.testing.assert_close(out, hclip)

        transforms.RandomHorizontalFlipVideo().__repr__()
//  [python]
def test_with_example_cases_multi_tasks_k_None(self) -> None:
        input = torch.tensor([[0.1, 0.5, 0.9]]).repeat(7, 1)
        target = torch.tensor(
            [
                [0, 0, 1],
                [0, 1, 0],
                [0, 1, 1],
                [1, 0, 0],
                [1, 0, 1],
                [1, 1, 0],
                [1, 1, 1],
            ]
        )
        expected_rp = torch.tensor([1 / 1, 1 / 1, 2 / 2, 0 / 2, 1 / 2, 1 / 2, 2 / 3])
        t.assert_close(
            retrieval_recall(input, target, k=2, limit_k_to_size=False, num_tasks=7),
            expected_rp,
        )
//  [python]
def test_with_example_cases_multi_tasks_k_None(self) -> None:
        input = torch.tensor([[0.1, 0.5, 0.9]]).repeat(7, 1)
        target = torch.tensor(
            [
                [0, 0, 1],
                [0, 1, 0],
                [0, 1, 1],
                [1, 0, 0],
                [1, 0, 1],
                [1, 1, 0],
                [1, 1, 1],
            ]
        )
        expected_rp = torch.tensor([1 / 1, 1 / 1, 2 / 2, 0 / 2, 1 / 2, 1 / 2, 2 / 3])
        torch.testing.assert_close(
            retrieval_recall(input, target, k=2, limit_k_to_size=False, num_tasks=7),
            expected_rp,
        )
//  [python]
def linear_test(self, Module, profiler_output_path=""):
        D_in = 10
        H = 5
        D_out = 15
        B = 8
        NUM_INPUTS = 2

        module = Module(D_in, H, D_out)

        inputs = []

        for i in range(NUM_INPUTS):
            inputs.append([torch.randn(B, D_in), torch.randn(B, D_in)])
        bench = ThroughputBenchmark(module)

        for input in inputs:
            # can do both args and kwargs here
            bench.add_input(input[0], x2=input[1])

        for i in range(NUM_INPUTS):
            # or just unpack the list of inputs
            module_result = module(*inputs[i])
            bench_result = bench.run_once(*inputs[i])
            torch.testing.assert_allclose(bench_result, module_result)

        stats = bench.benchmark(
            num_calling_threads=4,
            num_warmup_iters=100,
            num_iters=1000,
            profiler_output_path=profiler_output_path,
        )

        print(stats)
//  [python]
def linear_test(self, Module, profiler_output_path=""):
        D_in = 10
        H = 5
        D_out = 15
        B = 8
        NUM_INPUTS = 2

        module = Module(D_in, H, D_out)

        inputs = []

        for i in range(NUM_INPUTS):
            inputs.append([torch.randn(B, D_in), torch.randn(B, D_in)])
        bench = ThroughputBenchmark(module)

        for input in inputs:
            # can do both args and kwargs here
            bench.add_input(input[0], x2=input[1])

        for i in range(NUM_INPUTS):
            # or just unpack the list of inputs
            module_result = module(*inputs[i])
            bench_result = bench.run_once(*inputs[i])
            torch.testing.assert_close(bench_result, module_result)

        stats = bench.benchmark(
            num_calling_threads=4,
            num_warmup_iters=100,
            num_iters=1000,
            profiler_output_path=profiler_output_path,
        )

        print(stats)
//  [python]
def _test_identity(
        self,
        shape,
        elementwise,
        dtype="float16",
        test_name="identity",
    ) -> None:
        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )
        Y = ops.identity()(X)
        if elementwise:
            Y = ops.elementwise(FuncEnum.ADD)(X, Y)
        Y._attrs["name"] = "output"
        Y._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self.assertEqual(len(module.debug_sorted_graph), 2)
        self._test_id += 1

        x_pt = get_random_torch_tensor(shape, dtype=dtype)
        if elementwise:
            y_pt = 2 * x_pt
        else:
            y_pt = x_pt

        y = torch.empty_like(y_pt)

        module.run_with_tensors([x_pt], [y])
        torch.testing.assert_allclose(y, y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_identity(
        self,
        shape,
        elementwise,
        dtype="float16",
        test_name="identity",
    ) -> None:
        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )
        Y = ops.identity()(X)
        if elementwise:
            Y = ops.elementwise(FuncEnum.ADD)(X, Y)
        Y._attrs["name"] = "output"
        Y._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self.assertEqual(len(module.debug_sorted_graph), 2)
        self._test_id += 1

        x_pt = get_random_torch_tensor(shape, dtype=dtype)
        if elementwise:
            y_pt = 2 * x_pt
        else:
            y_pt = x_pt

        y = torch.empty_like(y_pt)

        module.run_with_tensors([x_pt], [y])
        torch.testing.assert_close(y, y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_with_example_cases_one_task(self) -> None:
        input = torch.tensor([0.1, 0.4, 0.6, 0.2, 0.5, 0.7, 0.3])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        t.assert_close(
            retrieval_recall(
                input,
                target,
            ),
            torch.tensor(4 / 4),
        )
        t.assert_close(
            retrieval_recall(
                input,
                target,
                k=2,
            ),
            torch.tensor(1 / 4),
        )
        t.assert_close(
            retrieval_recall(
                input,
                target,
                k=3,
            ),
            torch.tensor(2 / 4),
        )
        t.assert_close(
            retrieval_recall(
                input,
                target,
                k=4,
            ),
            torch.tensor(2 / 4),
        )
        t.assert_close(
            retrieval_recall(
                input,
                target,
                k=5,
            ),
            torch.tensor(3 / 4),
        )
//  [python]
def test_with_example_cases_one_task(self) -> None:
        input = torch.tensor([0.1, 0.4, 0.6, 0.2, 0.5, 0.7, 0.3])
        target = torch.tensor([0, 0, 1, 1, 1, 0, 1])

        torch.testing.assert_close(
            retrieval_recall(
                input,
                target,
            ),
            torch.tensor(4 / 4),
        )
        torch.testing.assert_close(
            retrieval_recall(
                input,
                target,
                k=2,
            ),
            torch.tensor(1 / 4),
        )
        torch.testing.assert_close(
            retrieval_recall(
                input,
                target,
                k=3,
            ),
            torch.tensor(2 / 4),
        )
        torch.testing.assert_close(
            retrieval_recall(
                input,
                target,
                k=4,
            ),
            torch.tensor(2 / 4),
        )
        torch.testing.assert_close(
            retrieval_recall(
                input,
                target,
                k=5,
            ),
            torch.tensor(3 / 4),
        )
//  [python]
def test_clone_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )
    original_feature_net_dict = dict(evolvable_mlp.feature_net.named_parameters())
    original_value_net_dict = dict(evolvable_mlp.value_net.named_parameters())
    original_advantage_net_dict = dict(evolvable_mlp.advantage_net.named_parameters())
    clone = evolvable_mlp.clone()
    clone_feature_net = clone.feature_net
    clone_value_net = clone.value_net
    clone_advantage_net = clone.advantage_net
    assert isinstance(clone, EvolvableMLP)
    print(clone.init_dict)
    print(evolvable_mlp.init_dict)
    for key, value in clone.init_dict.items():
        assert key in evolvable_mlp.init_dict
        if key == "support":
            torch.testing.assert_allclose(evolvable_mlp.init_dict[key], value)
        else:
            assert evolvable_mlp.init_dict[key] == value
    assert str(clone.state_dict()) == str(evolvable_mlp.state_dict())
    for key, param in clone_feature_net.named_parameters():
        torch.testing.assert_allclose(param, original_feature_net_dict[key])
    for key, param in clone_value_net.named_parameters():
        torch.testing.assert_allclose(param, original_value_net_dict[key])
    for key, param in clone_advantage_net.named_parameters():
        torch.testing.assert_allclose(param, original_advantage_net_dict[key])
//  [python]
def test_clone_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )
    original_feature_net_dict = dict(evolvable_mlp.feature_net.named_parameters())
    original_value_net_dict = dict(evolvable_mlp.value_net.named_parameters())
    original_advantage_net_dict = dict(evolvable_mlp.advantage_net.named_parameters())
    clone = evolvable_mlp.clone()
    clone_feature_net = clone.feature_net
    clone_value_net = clone.value_net
    clone_advantage_net = clone.advantage_net
    assert isinstance(clone, EvolvableMLP)
    print(clone.init_dict)
    print(evolvable_mlp.init_dict)
    for key, value in clone.init_dict.items():
        assert key in evolvable_mlp.init_dict
        if key == "support":
            torch.testing.assert_close(evolvable_mlp.init_dict[key], value)
        else:
            assert evolvable_mlp.init_dict[key] == value
    assert str(clone.state_dict()) == str(evolvable_mlp.state_dict())
    for key, param in clone_feature_net.named_parameters():
        torch.testing.assert_close(param, original_feature_net_dict[key])
    for key, param in clone_value_net.named_parameters():
        torch.testing.assert_close(param, original_value_net_dict[key])
    for key, param in clone_advantage_net.named_parameters():
        torch.testing.assert_close(param, original_advantage_net_dict[key])
//  [python]
def test_custom_joint_pass_pre(self):
        with config.patch(joint_custom_pre_pass=change_cos_pass):

            def g(x):
                return x.sin().sin().sin()

            def f(x):
                return x.cos().cos().cos()

            x = torch.randn(8, dtype=torch.float32)
            torch.testing.assert_allclose(torch.compile(f)(x), g(x))
//  [python]
def test_custom_joint_pass_pre(self):
        with config.patch(joint_custom_pre_pass=change_cos_pass):

            def g(x):
                return x.sin().sin().sin()

            def f(x):
                return x.cos().cos().cos()

            x = torch.randn(8, dtype=torch.float32)
            torch.testing.assert_close(torch.compile(f)(x), g(x))
//  [python]
def compare_jax_to_torch(
    jax_outputs: Any,
    torch_outputs: Any,
    tols: dict[str, float] | None = None,
) -> None:
  """Compares numerically Jax values to PyTorch values."""
  tols = tols or {}
  jax_leaves = jax.tree_util.tree_leaves(jax_outputs)
  torch_leaves = jax.tree_util.tree_leaves(
      torch_outputs,
      is_leaf=lambda x: isinstance(x, torch.Tensor),
  )
  assert len(jax_outputs) == len(torch_outputs)

  for jax_array, torch_array in zip(jax_leaves, torch_leaves):
    converted_array = conversion.jax_array_to_torch_tensor(jax_array)
    t.assert_close(torch_array, converted_array, **tols)
//  [python]
def compare_jax_to_torch(
    jax_outputs: Any,
    torch_outputs: Any,
    tols: dict[str, float] | None = None,
) -> None:
  """Compares numerically Jax values to PyTorch values."""
  tols = tols or {}
  jax_leaves = jax.tree_util.tree_leaves(jax_outputs)
  torch_leaves = jax.tree_util.tree_leaves(
      torch_outputs,
      is_leaf=lambda x: isinstance(x, torch.Tensor),
  )
  assert len(jax_outputs) == len(torch_outputs)

  for jax_array, torch_array in zip(jax_leaves, torch_leaves):
    converted_array = conversion.jax_array_to_torch_tensor(jax_array)
    torch.testing.assert_close(torch_array, converted_array, **tols)
//  [python]
def test_worker_prepare_inputs_for_prompt():
    worker = Worker(None, None, None)
    worker.block_size = 16
    batch_size = random.randint(1, 256)
    prompt_lens = []
    seq_group_metadata_list = []
    for i in range(batch_size):
        # make sure all tokens fit into one block
        prompt_len = i % (worker.block_size - 1) + 1
        prompt_lens.append(prompt_len)
        seq_data = list(range(prompt_len))
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData(seq_data)},
                sampling_params=SamplingParams(temperature=0),
                block_tables={0: [1]},
            ))
    expected_selected_token_indices = []
    selected_token_start_idx = 0
    max_seq_len = max(prompt_lens)
    for prompt_len in prompt_lens:
        expected_selected_token_indices.append(selected_token_start_idx +
                                               prompt_len - 1)
        selected_token_start_idx += max_seq_len
    input_tokens, input_positions, input_metadata = worker._prepare_inputs(
        seq_group_metadata_list)
    assert input_tokens.shape == input_positions.shape == (batch_size,
                                                           max_seq_len)
    t.assert_allclose(input_tokens, input_positions)
    actual = input_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    t.assert_allclose(actual, expected)
//  [python]
def test_worker_prepare_inputs_for_prompt():
    worker = Worker(None, None, None)
    worker.block_size = 16
    batch_size = random.randint(1, 256)
    prompt_lens = []
    seq_group_metadata_list = []
    for i in range(batch_size):
        # make sure all tokens fit into one block
        prompt_len = i % (worker.block_size - 1) + 1
        prompt_lens.append(prompt_len)
        seq_data = list(range(prompt_len))
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData(seq_data)},
                sampling_params=SamplingParams(temperature=0),
                block_tables={0: [1]},
            ))
    expected_selected_token_indices = []
    selected_token_start_idx = 0
    max_seq_len = max(prompt_lens)
    for prompt_len in prompt_lens:
        expected_selected_token_indices.append(selected_token_start_idx +
                                               prompt_len - 1)
        selected_token_start_idx += max_seq_len
    input_tokens, input_positions, input_metadata = worker._prepare_inputs(
        seq_group_metadata_list)
    assert input_tokens.shape == input_positions.shape == (batch_size,
                                                           max_seq_len)
    torch.testing.assert_close(input_tokens, input_positions)
    actual = input_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    torch.testing.assert_close(actual, expected)
//  [python]
def _test_common(
        self,
        mod,
        inputs,
        matcher_count,
        matcher_nodes,
        atol=1e-5,
        rtol=1.3e-6,
    ):
        counters.clear()
        maybe_autocast = contextlib.nullcontext()
        with torch.no_grad(), maybe_autocast:
            clone_inputs = self._clone_inputs(inputs)
            expected = mod(*inputs)
            actual = torch.compile(mod)(*clone_inputs)
            self.assertEqual(actual,expected)
            self.assertEqual(
                counters["inductor"]["pattern_matcher_count"], matcher_count
            )
            self.assertEqual(
                counters["inductor"]["pattern_matcher_nodes"],
                matcher_nodes,
            )
//  [python]
def _test_common(
        self,
        mod,
        inputs,
        matcher_count,
        matcher_nodes,
        atol=1e-5,
        rtol=1.3e-6,
    ):
        counters.clear()
        maybe_autocast = contextlib.nullcontext()
        with torch.no_grad(), maybe_autocast:
            clone_inputs = self._clone_inputs(inputs)
            expected = mod(*inputs)
            actual = torch.compile(mod)(*clone_inputs)
            torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
            self.assertEqual(
                counters["inductor"]["pattern_matcher_count"], matcher_count
            )
            self.assertEqual(
                counters["inductor"]["pattern_matcher_nodes"],
                matcher_nodes,
            )
//  [python]
def test_custom_joint_pass_post(self):
        with config.patch(joint_custom_post_pass=change_cos_pass):

            def g(x):
                return x.sin().sin().sin()

            def f(x):
                return x.cos().cos().cos()

            x = torch.randn(8, dtype=torch.float32)
            torch.testing.assert_allclose(torch.compile(f)(x), g(x))
//  [python]
def test_custom_joint_pass_post(self):
        with config.patch(joint_custom_post_pass=change_cos_pass):

            def g(x):
                return x.sin().sin().sin()

            def f(x):
                return x.cos().cos().cos()

            x = torch.randn(8, dtype=torch.float32)
            torch.testing.assert_close(torch.compile(f)(x), g(x))
//  [python]
def test_soft_round_conditional_mean():
    for offset in range(-5, 5):
        x = torch.linspace(offset + 0.001, offset + 0.999, 100)

        t.assert_allclose(
            torch.round(x),
            soft_round_conditional_mean(x, alpha=5000.0),
            atol=0.001,
            rtol=0.001,
        )
//  [python]
def test_soft_round_conditional_mean():
    for offset in range(-5, 5):
        x = torch.linspace(offset + 0.001, offset + 0.999, 100)

        torch.testing.assert_close(
            torch.round(x),
            soft_round_conditional_mean(x, alpha=5000.0),
            atol=0.001,
            rtol=0.001,
        )
//  [python]
def test_rematerialize_from_meta(self) -> None:
        pw = PositionWeightedModule(max_feature_length=10, device=torch.device("meta"))
        self.assertTrue(pw.position_weight.is_meta)

        # Re-materialize on cuda
        init_parameters(pw, torch.device("cuda"))
        self.assertTrue(not pw.position_weight.is_meta)
        self.assertEqual(pw.position_weight, torch.ones_like(pw.position_weight))
//  [python]
def test_rematerialize_from_meta(self) -> None:
        pw = PositionWeightedModule(max_feature_length=10, device=torch.device("meta"))
        self.assertTrue(pw.position_weight.is_meta)

        # Re-materialize on cuda
        init_parameters(pw, torch.device("cuda"))
        self.assertTrue(not pw.position_weight.is_meta)
        torch.testing.assert_close(
            pw.position_weight, torch.ones_like(pw.position_weight)
        )
//  [python]
def test_rematerialize_from_meta(self) -> None:
        pwmc = PositionWeightedModuleCollection(
            max_feature_lengths={"f1": 10, "f2": 10},
            device=torch.device("meta"),
        )
        self.assertTrue(all(param.is_meta for param in pwmc.position_weights.values()))

        # Re-materialize on cuda
        init_parameters(pwmc, torch.device("cuda"))
        for key, param in pwmc.position_weights.items():
            self.assertTrue(not param.is_meta)
            self.assertTrue(pwmc.position_weights_dict[key] is param)
            self.assertEqual(param, torch.ones_like(param))
//  [python]
def test_rematerialize_from_meta(self) -> None:
        pwmc = PositionWeightedModuleCollection(
            max_feature_lengths={"f1": 10, "f2": 10},
            device=torch.device("meta"),
        )
        self.assertTrue(all(param.is_meta for param in pwmc.position_weights.values()))

        # Re-materialize on cuda
        init_parameters(pwmc, torch.device("cuda"))
        for key, param in pwmc.position_weights.items():
            self.assertTrue(not param.is_meta)
            self.assertTrue(pwmc.position_weights_dict[key] is param)
            torch.testing.assert_close(param, torch.ones_like(param))
//  [python]
def _check_ddp_fsdp_param_parity(self, ddp_model: DDP, fsdp_model: FSDP):
        with FSDP.summon_full_params(fsdp_model):
            for (n1, p1), (n2, p2) in zip(
                ddp_model.module.named_parameters(), fsdp_model.named_parameters()
            ):
                # Allow for FSDP prefixes
                self.assertEqual(n1, clean_tensor_name(n2))
                self.assertEqual(p1.tolist(),p2.tolist())
//  [python]
def _check_ddp_fsdp_param_parity(self, ddp_model: DDP, fsdp_model: FSDP):
        with FSDP.summon_full_params(fsdp_model):
            for (n1, p1), (n2, p2) in zip(
                ddp_model.module.named_parameters(), fsdp_model.named_parameters()
            ):
                # Allow for FSDP prefixes
                self.assertEqual(n1, clean_tensor_name(n2))
                torch.testing.assert_close(p1, p2)
//  [python]
def test_mm_backward(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        bsz = 4
        k = 5
        n = 7
        shard_spec = [Shard(1)]
        test_lhs = torch.rand((bsz, 3 * self.world_size, k), device=self.device_type, requires_grad=True)
        test_rhs = torch.rand((k, n), device=self.device_type, requires_grad=True)
        base_lhs = deepcopy(test_lhs)
        base_rhs = deepcopy(test_rhs)

        # DTensor operation
        d_lhs = distribute_tensor(test_lhs, device_mesh, shard_spec)
        d_rhs = distribute_tensor(test_rhs, device_mesh, [Replicate()])
        d_reshaped_lhs = d_lhs.reshape((-1, k))
        d_out = torch.mm(d_reshaped_lhs, d_rhs)
        d_out = d_out.reshape((bsz, -1, n)).redistribute(device_mesh, [Replicate()])
        d_out.sum().to_local().backward()

        # Tensor operation
        base_out = torch.mm(base_lhs.reshape((-1, k)), base_rhs)
        base_out.reshape((bsz, -1, n)).sum().backward()

        # compare grad of lhs, note that dtensor's grad is Sharded
        torch.testing.assert_allclose(d_lhs.grad.redistribute(device_mesh, [Replicate()])._local_tensor, base_lhs.grad)
        # compare grad of rhs, note that dtensor's grad is Partial
        torch.testing.assert_allclose(d_rhs.grad.redistribute(device_mesh, [Replicate()])._local_tensor, base_rhs.grad)
//  [python]
def test_mm_backward(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        bsz = 4
        k = 5
        n = 7
        shard_spec = [Shard(1)]
        test_lhs = torch.rand((bsz, 3 * self.world_size, k), device=self.device_type, requires_grad=True)
        test_rhs = torch.rand((k, n), device=self.device_type, requires_grad=True)
        base_lhs = deepcopy(test_lhs)
        base_rhs = deepcopy(test_rhs)

        # DTensor operation
        d_lhs = distribute_tensor(test_lhs, device_mesh, shard_spec)
        d_rhs = distribute_tensor(test_rhs, device_mesh, [Replicate()])
        d_reshaped_lhs = d_lhs.reshape((-1, k))
        d_out = torch.mm(d_reshaped_lhs, d_rhs)
        d_out = d_out.reshape((bsz, -1, n)).redistribute(device_mesh, [Replicate()])
        d_out.sum().to_local().backward()

        # Tensor operation
        base_out = torch.mm(base_lhs.reshape((-1, k)), base_rhs)
        base_out.reshape((bsz, -1, n)).sum().backward()

        # compare grad of lhs, note that dtensor's grad is Sharded
        torch.testing.assert_close(d_lhs.grad.redistribute(device_mesh, [Replicate()])._local_tensor, base_lhs.grad)
        # compare grad of rhs, note that dtensor's grad is Partial
        torch.testing.assert_close(d_rhs.grad.redistribute(device_mesh, [Replicate()])._local_tensor, base_rhs.grad)
//  [python]
def test_populate_weights(self) -> None:
        pw = PositionWeightedModule(max_feature_length=10)

        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )
        features = features.to_dict()

        jt = features["f1"]
        weighted_features = pw(jt)

        self.assertEqual(weighted_features.weights().size(), (3,))

        pw_f1_ref = torch.gather(
            pw.state_dict()["position_weight"], 0, torch.tensor([0, 1, 0])
        )

        pw_f1 = weighted_features.weights().detach()
        self.assertTrue(torch.allclose(pw_f1_ref, pw_f1))

        position_weighted_module_gm = symbolic_trace(pw)
        position_weighted_module_gm_script = torch.jit.script(
            position_weighted_module_gm
        )

        weighted_features_gm_script = position_weighted_module_gm_script(jt)
        torch.testing.assert_allclose(
            weighted_features.values(), weighted_features_gm_script.values()
        )
        torch.testing.assert_allclose(
            weighted_features.lengths(), weighted_features_gm_script.lengths()
        )
//  [python]
def test_populate_weights(self) -> None:
        pw = PositionWeightedModule(max_feature_length=10)

        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )
        features = features.to_dict()

        jt = features["f1"]
        weighted_features = pw(jt)

        self.assertEqual(weighted_features.weights().size(), (3,))

        pw_f1_ref = torch.gather(
            pw.state_dict()["position_weight"], 0, torch.tensor([0, 1, 0])
        )

        pw_f1 = weighted_features.weights().detach()
        self.assertTrue(torch.allclose(pw_f1_ref, pw_f1))

        position_weighted_module_gm = symbolic_trace(pw)
        position_weighted_module_gm_script = torch.jit.script(
            position_weighted_module_gm
        )

        weighted_features_gm_script = position_weighted_module_gm_script(jt)
        torch.testing.assert_close(
            weighted_features.values(), weighted_features_gm_script.values()
        )
        torch.testing.assert_close(
            weighted_features.lengths(), weighted_features_gm_script.lengths()
        )
//  [python]
def _check_fsdp_parameter_parity(self, fsdp1: FSDP, fsdp2: FSDP) -> None:
        """Checks that two FSDP instances have the same model parameters."""
        with FSDP.summon_full_params(fsdp1), FSDP.summon_full_params(fsdp2):
            for (n1, p1), (n2, p2) in zip(
                fsdp1.named_parameters(),
                fsdp2.named_parameters(),
            ):
                self.assertEqual(n1, n2)
                self.assertEqual(p1.tolist(),p2.tolist())
//  [python]
def _check_fsdp_parameter_parity(self, fsdp1: FSDP, fsdp2: FSDP) -> None:
        """Checks that two FSDP instances have the same model parameters."""
        with FSDP.summon_full_params(fsdp1), FSDP.summon_full_params(fsdp2):
            for (n1, p1), (n2, p2) in zip(
                fsdp1.named_parameters(),
                fsdp2.named_parameters(),
            ):
                self.assertEqual(n1, n2)
                torch.testing.assert_close(p1, p2)
//  [python]
def _check_train_parity(
        self,
        ddp_model: DDP,
        ddp_optim: torch.optim.Optimizer,
        fsdp_model: FSDP,
        fsdp_optim: torch.optim.Optimizer,
        set_to_none: bool,
        num_iters: int = 10,
    ):
        """Checks training parity between DDP and FSDP."""
        device = torch.device("cuda")
        for i in range(num_iters):
            iter_losses = []
            for model, optim in ((ddp_model, ddp_optim), (fsdp_model, fsdp_optim)):
                module = model.module
                # Test two different `zero_grad()` timings
                if i % 2 == 0:
                    optim.zero_grad(set_to_none=set_to_none)  # pre-forward
                inp = module.get_input(device)
                output = model(*inp)
                loss = module.get_loss(inp, output).to(device)
                iter_losses.append(loss)
                if i % 2 == 1:
                    optim.zero_grad(set_to_none=set_to_none)  # pre-backward
                module.run_backward(loss)
                # Perform the DDP optimizer step on CPU to match FSDP if needed
                if model is ddp_model and fsdp_model.cpu_offload.offload_params:
                    model.to(torch.device("cpu"))
                optim.step()
                if model is ddp_model and fsdp_model.cpu_offload.offload_params:
                    model.to(device)
            torch.testing.assert_allclose(iter_losses[0], iter_losses[1])
            iter_losses.clear()
        self._check_ddp_fsdp_param_parity(ddp_model, fsdp_model)
//  [python]
def _check_train_parity(
        self,
        ddp_model: DDP,
        ddp_optim: torch.optim.Optimizer,
        fsdp_model: FSDP,
        fsdp_optim: torch.optim.Optimizer,
        set_to_none: bool,
        num_iters: int = 10,
    ):
        """Checks training parity between DDP and FSDP."""
        device = torch.device("cuda")
        for i in range(num_iters):
            iter_losses = []
            for model, optim in ((ddp_model, ddp_optim), (fsdp_model, fsdp_optim)):
                module = model.module
                # Test two different `zero_grad()` timings
                if i % 2 == 0:
                    optim.zero_grad(set_to_none=set_to_none)  # pre-forward
                inp = module.get_input(device)
                output = model(*inp)
                loss = module.get_loss(inp, output).to(device)
                iter_losses.append(loss)
                if i % 2 == 1:
                    optim.zero_grad(set_to_none=set_to_none)  # pre-backward
                module.run_backward(loss)
                # Perform the DDP optimizer step on CPU to match FSDP if needed
                if model is ddp_model and fsdp_model.cpu_offload.offload_params:
                    model.to(torch.device("cpu"))
                optim.step()
                if model is ddp_model and fsdp_model.cpu_offload.offload_params:
                    model.to(device)
            torch.testing.assert_close(iter_losses[0], iter_losses[1])
            iter_losses.clear()
        self._check_ddp_fsdp_param_parity(ddp_model, fsdp_model)
//  [python]
def _check_param_parity(self, ddp_model: DDP, fsdp_model: FSDP):
        with FSDP.summon_full_params(fsdp_model):
            for (n1, p1), (n2, p2) in zip(
                ddp_model.module.named_parameters(),
                fsdp_model.named_parameters(),
            ):
                self.assertEqual(n1, n2)
                self.assertEqual(p1.dtype, p2.dtype)
                self.assertTrue(torch.allclose(p1, p2))
//  [python]
def _check_param_parity(self, ddp_model: DDP, fsdp_model: FSDP):
        with FSDP.summon_full_params(fsdp_model):
            for (n1, p1), (n2, p2) in zip(
                ddp_model.module.named_parameters(),
                fsdp_model.named_parameters(),
            ):
                self.assertEqual(n1, n2)
                torch.testing.assert_close(p1, p2)
//  [python]
def test_populate_weights(self) -> None:
        position_weighted_module_collection = PositionWeightedModuleCollection(
            {"f1": 10, "f2": 10}
        )

        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )

        fp_kjt = position_weighted_module_collection(features)

        position_weighted_module_collection_gm = symbolic_trace(
            position_weighted_module_collection
        )
        position_weighted_module_collection_gm_script = torch.jit.script(
            position_weighted_module_collection_gm
        )
        fp_kjt_gm_script = position_weighted_module_collection_gm_script(features)

        t.assert_close(fp_kjt.values(), fp_kjt_gm_script.values())
        t.assert_close(fp_kjt.lengths(), fp_kjt_gm_script.lengths())
        t.assert_close(
            fp_kjt.length_per_key(), fp_kjt_gm_script.length_per_key()
        )

        empty_kjt = KeyedJaggedTensor.from_lengths_sync(
            keys=[],
            values=torch.tensor([], dtype=torch.int32),
            lengths=torch.tensor([], dtype=torch.int32),
        )

        empty_fp_kjt = position_weighted_module_collection(empty_kjt)
        empty_fp_kjt_gm_script = position_weighted_module_collection_gm_script(
            empty_kjt
        )

        t.assert_close(
            empty_fp_kjt.values(), empty_fp_kjt_gm_script.values()
        )
        t.assert_close(
            empty_fp_kjt.lengths(), empty_fp_kjt_gm_script.lengths()
        )
        t.assert_close(
            empty_fp_kjt.length_per_key(), empty_fp_kjt_gm_script.length_per_key()
        )
//  [python]
def test_populate_weights(self) -> None:
        position_weighted_module_collection = PositionWeightedModuleCollection(
            {"f1": 10, "f2": 10}
        )

        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )

        fp_kjt = position_weighted_module_collection(features)

        position_weighted_module_collection_gm = symbolic_trace(
            position_weighted_module_collection
        )
        position_weighted_module_collection_gm_script = torch.jit.script(
            position_weighted_module_collection_gm
        )
        fp_kjt_gm_script = position_weighted_module_collection_gm_script(features)

        torch.testing.assert_close(fp_kjt.values(), fp_kjt_gm_script.values())
        torch.testing.assert_close(fp_kjt.lengths(), fp_kjt_gm_script.lengths())
        torch.testing.assert_close(
            fp_kjt.length_per_key(), fp_kjt_gm_script.length_per_key()
        )

        empty_kjt = KeyedJaggedTensor.from_lengths_sync(
            keys=[],
            values=torch.tensor([], dtype=torch.int32),
            lengths=torch.tensor([], dtype=torch.int32),
        )

        empty_fp_kjt = position_weighted_module_collection(empty_kjt)
        empty_fp_kjt_gm_script = position_weighted_module_collection_gm_script(
            empty_kjt
        )

        torch.testing.assert_close(
            empty_fp_kjt.values(), empty_fp_kjt_gm_script.values()
        )
        torch.testing.assert_close(
            empty_fp_kjt.lengths(), empty_fp_kjt_gm_script.lengths()
        )
        torch.testing.assert_close(
            empty_fp_kjt.length_per_key(), empty_fp_kjt_gm_script.length_per_key()
        )
//  [python]
def _test_gemm_rcr_fast_gelu(
        self,
        Ms,
        test_name,
        K=1024,
        N=64,
        dtype="float16",
    ):
        MDim = shape_utils.gen_int_var_min_max(Ms, name="m")
        X = Tensor(
            shape=[MDim, IntImm(K)],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(
            shape=[IntImm(N), IntImm(K)],
            dtype=dtype,
            name="input_1",
            is_input=True,
        )

        OP = ops.gemm_rcr_fast_gelu()

        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        module = compile_model(
            Y,
            detect_target(),
            "./tmp",
            f"gemm_rcr_fast_gelu_{test_name}",
        )

        for M in Ms:
            logging.info(f"Testing {M=}")

            X_pt = get_random_torch_tensor([M, K], dtype)
            W_pt = get_random_torch_tensor([N, K], dtype)
            Y_pt = NewGELUActivation()(torch.nn.functional.linear(X_pt, W_pt))
            y = get_torch_empty_tensor([M, N], dtype)
            module.run_with_tensors(
                {"input_0": X_pt, "input_1": W_pt},
                [y],
            )
            t.assert_allclose(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def _test_gemm_rcr_fast_gelu(
        self,
        Ms,
        test_name,
        K=1024,
        N=64,
        dtype="float16",
    ):
        MDim = shape_utils.gen_int_var_min_max(Ms, name="m")
        X = Tensor(
            shape=[MDim, IntImm(K)],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(
            shape=[IntImm(N), IntImm(K)],
            dtype=dtype,
            name="input_1",
            is_input=True,
        )

        OP = ops.gemm_rcr_fast_gelu()

        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        module = compile_model(
            Y,
            detect_target(),
            "./tmp",
            f"gemm_rcr_fast_gelu_{test_name}",
        )

        for M in Ms:
            logging.info(f"Testing {M=}")

            X_pt = get_random_torch_tensor([M, K], dtype)
            W_pt = get_random_torch_tensor([N, K], dtype)
            Y_pt = NewGELUActivation()(torch.nn.functional.linear(X_pt, W_pt))
            y = get_torch_empty_tensor([M, N], dtype)
            module.run_with_tensors(
                {"input_0": X_pt, "input_1": W_pt},
                [y],
            )
            torch.testing.assert_close(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def test_save_and_load_trt_module(self):
        class TestModule(torch.nn.Module):
            def forward(self, x):
                return x + x

        inputs = [torch.randn(1, 1)]
        mod = TestModule().eval()
        ref_output = mod(*inputs)

        mod = acc_tracer.trace(mod, inputs)
        interp = TRTInterpreter(mod, input_specs=InputTensorSpec.from_tensors(inputs))
        trt_mod = TRTModule(*interp.run(lower_precision=LowerPrecision.FP32))
        torch.save(trt_mod, "trt.pt")
        reload_trt_mod = torch.load("trt.pt")

        torch.testing.assert_allclose(
            reload_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
        )
        os.remove(f"{os.getcwd()}/trt.pt")
//  [python]
def test_save_and_load_trt_module(self):
        class TestModule(torch.nn.Module):
            def forward(self, x):
                return x + x

        inputs = [torch.randn(1, 1)]
        mod = TestModule().eval()
        ref_output = mod(*inputs)

        mod = acc_tracer.trace(mod, inputs)
        interp = TRTInterpreter(mod, input_specs=InputTensorSpec.from_tensors(inputs))
        trt_mod = TRTModule(*interp.run(lower_precision=LowerPrecision.FP32))
        torch.save(trt_mod, "trt.pt")
        reload_trt_mod = torch.load("trt.pt")

        torch.testing.assert_close(
            reload_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
        )
        os.remove(f"{os.getcwd()}/trt.pt")