jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_against_falcon_40b():
    ours_config = Config.from_name("falcon-40b", n_layer=2, n_head=8, n_query_groups=4, n_embd=32)
    theirs_config = FalconConfig(
        vocab_size=ours_config.padded_vocab_size,
        hidden_size=ours_config.n_embd,
        num_hidden_layers=ours_config.n_layer,
        num_attention_heads=ours_config.n_head,
        num_kv_heads=ours_config.n_query_groups,
        new_decoder_architecture=True,
        parallel_attn=ours_config.parallel_residual,
        bias=ours_config.bias,
    )

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_falcon("40b", theirs_state_dict, ours_state_dict)

    theirs_model = FalconForCausalLM(theirs_config)
    # assign must be set to True for torch.testing.assert_close to pass
    theirs_model.load_state_dict(theirs_state_dict, assign=True)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    t.assert_close(ours_y, theirs_y)
//  [python]
def test_against_falcon_40b():
    ours_config = Config.from_name("falcon-40b", n_layer=2, n_head=8, n_query_groups=4, n_embd=32)
    theirs_config = FalconConfig(
        vocab_size=ours_config.padded_vocab_size,
        hidden_size=ours_config.n_embd,
        num_hidden_layers=ours_config.n_layer,
        num_attention_heads=ours_config.n_head,
        num_kv_heads=ours_config.n_query_groups,
        new_decoder_architecture=True,
        parallel_attn=ours_config.parallel_residual,
        bias=ours_config.bias,
    )

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_falcon("40b", theirs_state_dict, ours_state_dict)

    theirs_model = FalconForCausalLM(theirs_config)
    # assign must be set to True for torch.testing.assert_close to pass
    theirs_model.load_state_dict(theirs_state_dict, assign=True)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_chunked_cross_entropy(ignore_index, B):
    V = 50
    T = 25
    regular_logits = torch.randn(B, T, V)
    targets = torch.randint(0, V, (B, T))

    if ignore_index is not None:
        targets[:, [1, 4, 10, 19]] = ignore_index

    baseline_loss = F.cross_entropy(
        regular_logits.reshape(-1, regular_logits.size(-1)),
        targets.reshape(-1),
        ignore_index=(ignore_index if ignore_index is not None else -100),
    )

    ignore_index = ignore_index if ignore_index is not None else -100
    regular_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=0, ignore_index=ignore_index)
    assert torch.equal(baseline_loss, regular_loss)
    assert regular_loss.numel() == 1

    chunked_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=10, ignore_index=ignore_index)
    torch.testing.assert_allclose(chunked_loss, regular_loss)
    torch.testing.assert_allclose(chunked_loss, baseline_loss)

    logit_chunk_size = 6
    assert T % logit_chunk_size != 0  # ensure leftover
    chunked_logits = list(regular_logits.split(logit_chunk_size, dim=1))
    chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=0, ignore_index=ignore_index)
    torch.testing.assert_allclose(chunked_loss, regular_loss)
    torch.testing.assert_allclose(chunked_loss, baseline_loss)

    chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=10, ignore_index=ignore_index)
    torch.testing.assert_allclose(chunked_loss, regular_loss)
    torch.testing.assert_allclose(chunked_loss, baseline_loss)
//  [python]
def test_chunked_cross_entropy(ignore_index, B):
    V = 50
    T = 25
    regular_logits = torch.randn(B, T, V)
    targets = torch.randint(0, V, (B, T))

    if ignore_index is not None:
        targets[:, [1, 4, 10, 19]] = ignore_index

    baseline_loss = F.cross_entropy(
        regular_logits.reshape(-1, regular_logits.size(-1)),
        targets.reshape(-1),
        ignore_index=(ignore_index if ignore_index is not None else -100),
    )

    ignore_index = ignore_index if ignore_index is not None else -100
    regular_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=0, ignore_index=ignore_index)
    assert torch.equal(baseline_loss, regular_loss)
    assert regular_loss.numel() == 1

    chunked_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=10, ignore_index=ignore_index)
    torch.testing.assert_close(chunked_loss, regular_loss)
    torch.testing.assert_close(chunked_loss, baseline_loss)

    logit_chunk_size = 6
    assert T % logit_chunk_size != 0  # ensure leftover
    chunked_logits = list(regular_logits.split(logit_chunk_size, dim=1))
    chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=0, ignore_index=ignore_index)
    torch.testing.assert_close(chunked_loss, regular_loss)
    torch.testing.assert_close(chunked_loss, baseline_loss)

    chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=10, ignore_index=ignore_index)
    torch.testing.assert_close(chunked_loss, regular_loss)
    torch.testing.assert_close(chunked_loss, baseline_loss)
//  [python]
def test_against_original_open_llama_3b():
    ours_config = Config.from_name("open_llama_3b", n_layer=2, n_head=8, n_embd=32, intermediate_size=86)
    T = 5
    theirs_config = LlamaConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = LlamaForCausalLM(theirs_config)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_allclose(ours_y, theirs_y)
//  [python]
def test_against_original_open_llama_3b():
    ours_config = Config.from_name("open_llama_3b", n_layer=2, n_head=8, n_embd=32, intermediate_size=86)
    T = 5
    theirs_config = LlamaConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = LlamaForCausalLM(theirs_config)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_against_hf_llama2(ours_kwargs):
    ours_config = Config.from_name(
        padded_vocab_size=10000, n_layer=2, n_head=8, n_embd=32, intermediate_size=86, **ours_kwargs
    )
    T = 5
    theirs_config = LlamaConfig(
        vocab_size=ours_config.padded_vocab_size,
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
        rms_norm_eps=ours_config.norm_eps,
        num_query_value_heads=ours_config.n_query_groups,
        rope_theta=ours_config.rope_base,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = LlamaForCausalLM(theirs_config)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_allclose(ours_y, theirs_y)
//  [python]
def test_against_hf_llama2(ours_kwargs):
    ours_config = Config.from_name(
        padded_vocab_size=10000, n_layer=2, n_head=8, n_embd=32, intermediate_size=86, **ours_kwargs
    )
    T = 5
    theirs_config = LlamaConfig(
        vocab_size=ours_config.padded_vocab_size,
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
        rms_norm_eps=ours_config.norm_eps,
        num_query_value_heads=ours_config.n_query_groups,
        rope_theta=ours_config.rope_base,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = LlamaForCausalLM(theirs_config)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_against_original_stablelm_zephyr_3b():
    T = 5
    ours_config = Config.from_name("stablelm-zephyr-3b", n_layer=2, n_head=16, n_embd=32, intermediate_size=86)
    theirs_config = AutoConfig.from_pretrained(
        "stabilityai/stablelm-zephyr-3b",
        trust_remote_code=True,
        num_hidden_layers=ours_config.n_layer,
        num_attention_heads=ours_config.n_head,
        num_key_value_heads=ours_config.n_head,
        hidden_size=ours_config.n_embd,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = AutoModelForCausalLM.from_config(theirs_config, trust_remote_code=True)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_allclose(ours_y, theirs_y)
//  [python]
def test_against_original_stablelm_zephyr_3b():
    T = 5
    ours_config = Config.from_name("stablelm-zephyr-3b", n_layer=2, n_head=16, n_embd=32, intermediate_size=86)
    theirs_config = AutoConfig.from_pretrained(
        "stabilityai/stablelm-zephyr-3b",
        trust_remote_code=True,
        num_hidden_layers=ours_config.n_layer,
        num_attention_heads=ours_config.n_head,
        num_key_value_heads=ours_config.n_head,
        hidden_size=ours_config.n_embd,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    ours_model = GPT(ours_config)
    ours_state_dict = ours_model.state_dict()
    theirs_state_dict = {}
    copy_weights_llama(ours_config, theirs_state_dict, ours_state_dict)
    theirs_model = AutoModelForCausalLM.from_config(theirs_config, trust_remote_code=True)
    theirs_model.load_state_dict(theirs_state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_ellipsis_ops(self) -> None:
        x = torch.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
        for pattern in identity_patterns:
            self.assertEqual(rearrange(x, pattern), x, msg=pattern)

        for pattern1, pattern2 in equivalent_rearrange_patterns:
            self.assertEqual(
                rearrange(x, pattern1),
                rearrange(x, pattern2),
                msg=f"{pattern1} vs {pattern2}",
            )
//  [python]
def test_ellipsis_ops(self) -> None:
        x = torch.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
        for pattern in identity_patterns:
            torch.testing.assert_close(rearrange(x, pattern), x, msg=pattern)

        for pattern1, pattern2 in equivalent_rearrange_patterns:
            torch.testing.assert_close(
                rearrange(x, pattern1),
                rearrange(x, pattern2),
                msg=f"{pattern1} vs {pattern2}",
            )
//  [python]
def test_zero_grid_with_backed_symbols(self):
        class Repro(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x, b):
                return x + b

        example_inputs = (
            x := torch.randn((3, 2), device=self.device),
            torch.randn((1, 2), device=self.device),
        )
        torch._dynamo.mark_dynamic(x, index=0)  # Create dynamic symbol

        # Compile & run model where dynamic dim size > 0.
        so_path: str = AOTIRunnerUtil.compile(
            Repro(),
            example_inputs,
        )
        aot_inductor_module = AOTIRunnerUtil.load("cuda", so_path)
        aot_inductor_module(*example_inputs)

        # Re-run where dynamic dim size is 0.
        example_inputs = (
            torch.randn((0, 2), device=self.device),
            torch.randn((1, 2), device=self.device),
        )
        actual = aot_inductor_module(*example_inputs)
        expected = Repro()(*example_inputs)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_zero_grid_with_backed_symbols(self):
        class Repro(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x, b):
                return x + b

        example_inputs = (
            x := torch.randn((3, 2), device=self.device),
            torch.randn((1, 2), device=self.device),
        )
        torch._dynamo.mark_dynamic(x, index=0)  # Create dynamic symbol

        # Compile & run model where dynamic dim size > 0.
        so_path: str = AOTIRunnerUtil.compile(
            Repro(),
            example_inputs,
        )
        aot_inductor_module = AOTIRunnerUtil.load("cuda", so_path)
        aot_inductor_module(*example_inputs)

        # Re-run where dynamic dim size is 0.
        example_inputs = (
            torch.randn((0, 2), device=self.device),
            torch.randn((1, 2), device=self.device),
        )
        actual = aot_inductor_module(*example_inputs)
        expected = Repro()(*example_inputs)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_unsqueeze(self) -> None:
        x = torch.randn((2, 3, 4, 5))
        actual = rearrange(x, "b h w c -> b 1 h w 1 c")
        expected = x.unsqueeze(1).unsqueeze(-2)
        self.assertEqual(actual, expected)
//  [python]
def test_unsqueeze(self) -> None:
        x = torch.randn((2, 3, 4, 5))
        actual = rearrange(x, "b h w c -> b 1 h w 1 c")
        expected = x.unsqueeze(1).unsqueeze(-2)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_concatenations_and_stacking(self) -> None:
        for n_arrays in [1, 2, 5]:
            shapes: List[List[int]] = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
            for shape in shapes:
                arrays1 = [
                    torch.arange(i, i + np.prod(shape, dtype=int)).reshape(shape)
                    for i in range(n_arrays)
                ]
                result0 = torch.stack(arrays1)
                result1 = rearrange(arrays1, "...->...")
                t.assert_close(result0, result1)
//  [python]
def test_concatenations_and_stacking(self) -> None:
        for n_arrays in [1, 2, 5]:
            shapes: List[List[int]] = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
            for shape in shapes:
                arrays1 = [
                    torch.arange(i, i + np.prod(shape, dtype=int)).reshape(shape)
                    for i in range(n_arrays)
                ]
                result0 = torch.stack(arrays1)
                result1 = rearrange(arrays1, "...->...")
                torch.testing.assert_close(result0, result1)
//  [python]
def test_triton_kernel_reinterpret_view_mem_leak(self):
        # Check for memory leak when using user-defined Triton Kernel + AOTI.
        if self.device != "cuda":
            raise unittest.SkipTest("requires CUDA")

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

            def forward(self, x, y):
                out = torch.zeros_like(x)
                yy = y * y
                # reshape creates a ReinterpretView
                add_kernel[(4,)](x, yy.reshape_as(x), out, 4, 16)
                return out

        example_inputs = (
            torch.randn(4, 4, device="cuda"),
            torch.randn(1, 16, device="cuda"),
        )

        so_path: str = AOTIRunnerUtil.compile(
            Model(),
            example_inputs,
        )
        aot_inductor_module = AOTIRunnerUtil.load("cuda", so_path)

        # Don't assign outputs to a variable b/c it will allocate GPU memory.
        device: int = torch.cuda.current_device()
        mem_before = torch.cuda.memory_allocated(device)
        aot_inductor_module(*example_inputs)
        aot_inductor_module(*example_inputs)
        mem_after = torch.cuda.memory_allocated(device)
        self.assertEqual(mem_before, mem_after)

        actual = aot_inductor_module(*example_inputs)
        expected = Model()(*example_inputs)
        self.assertEqual(actual, expected)
//  [python]
def test_triton_kernel_reinterpret_view_mem_leak(self):
        # Check for memory leak when using user-defined Triton Kernel + AOTI.
        if self.device != "cuda":
            raise unittest.SkipTest("requires CUDA")

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

            def forward(self, x, y):
                out = torch.zeros_like(x)
                yy = y * y
                # reshape creates a ReinterpretView
                add_kernel[(4,)](x, yy.reshape_as(x), out, 4, 16)
                return out

        example_inputs = (
            torch.randn(4, 4, device="cuda"),
            torch.randn(1, 16, device="cuda"),
        )

        so_path: str = AOTIRunnerUtil.compile(
            Model(),
            example_inputs,
        )
        aot_inductor_module = AOTIRunnerUtil.load("cuda", so_path)

        # Don't assign outputs to a variable b/c it will allocate GPU memory.
        device: int = torch.cuda.current_device()
        mem_before = torch.cuda.memory_allocated(device)
        aot_inductor_module(*example_inputs)
        aot_inductor_module(*example_inputs)
        mem_after = torch.cuda.memory_allocated(device)
        self.assertEqual(mem_before, mem_after)

        actual = aot_inductor_module(*example_inputs)
        expected = Model()(*example_inputs)
        torch.testing.assert_close(actual, expected)
//  [python]
def test_squeeze(self) -> None:
        x = torch.randn((2, 1, 3, 4, 1, 5))
        actual = rearrange(x, "b 1 h w 1 c -> b h w c")
        expected = x.squeeze()
        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_squeeze(self) -> None:
        x = torch.randn((2, 1, 3, 4, 1, 5))
        actual = rearrange(x, "b 1 h w 1 c -> b h w c")
        expected = x.squeeze()
        torch.testing.assert_close(actual, expected)
//  [python]
def test_0_dim_tensor(self) -> None:
        x = expected = torch.tensor(1)
        actual = rearrange(x, "->")
        self.assertEqual(actual, expected)

        actual = rearrange(x, "... -> ...")
        self.assertEqual(actual, expected)
//  [python]
def test_0_dim_tensor(self) -> None:
        x = expected = torch.tensor(1)
        actual = rearrange(x, "->")
        torch.testing.assert_close(actual, expected)

        actual = rearrange(x, "... -> ...")
        torch.testing.assert_close(actual, expected)
//  [python]
def test_rearrange_permutations(self) -> None:
        # tests random permutation of axes against two independent numpy ways
        for n_axes in range(1, 10):
            input = torch.arange(2**n_axes).reshape([2] * n_axes)
            permutation = np.random.permutation(n_axes)
            left_expression = " ".join("i" + str(axis) for axis in range(n_axes))
            right_expression = " ".join("i" + str(axis) for axis in permutation)
            expression = left_expression + " -> " + right_expression
            result = rearrange(input, expression)

            for pick in np.random.randint(0, 2, [10, n_axes]):
                self.assertEqual(input[tuple(pick)], result[tuple(pick[permutation])])

        for n_axes in range(1, 10):
            input = torch.arange(2**n_axes).reshape([2] * n_axes)
            permutation = np.random.permutation(n_axes)
            left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1])
            right_expression = " ".join("i" + str(axis) for axis in permutation[::-1])
            expression = left_expression + " -> " + right_expression
            result = rearrange(input, expression)
            self.assertEqual(result.shape, input.shape)
            expected_result = torch.zeros_like(input)
            for original_axis, result_axis in enumerate(permutation):
                expected_result |= ((input >> original_axis) & 1) << result_axis

            self.assertEqual(result, expected_result)
//  [python]
def test_rearrange_permutations(self) -> None:
        # tests random permutation of axes against two independent numpy ways
        for n_axes in range(1, 10):
            input = torch.arange(2**n_axes).reshape([2] * n_axes)
            permutation = np.random.permutation(n_axes)
            left_expression = " ".join("i" + str(axis) for axis in range(n_axes))
            right_expression = " ".join("i" + str(axis) for axis in permutation)
            expression = left_expression + " -> " + right_expression
            result = rearrange(input, expression)

            for pick in np.random.randint(0, 2, [10, n_axes]):
                self.assertEqual(input[tuple(pick)], result[tuple(pick[permutation])])

        for n_axes in range(1, 10):
            input = torch.arange(2**n_axes).reshape([2] * n_axes)
            permutation = np.random.permutation(n_axes)
            left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1])
            right_expression = " ".join("i" + str(axis) for axis in permutation[::-1])
            expression = left_expression + " -> " + right_expression
            result = rearrange(input, expression)
            self.assertEqual(result.shape, input.shape)
            expected_result = torch.zeros_like(input)
            for original_axis, result_axis in enumerate(permutation):
                expected_result |= ((input >> original_axis) & 1) << result_axis

            torch.testing.assert_close(result, expected_result)
//  [python]
def test_rearrange_consistency(self) -> None:
        shape = [1, 2, 3, 5, 7, 11]
        x = torch.arange(int(np.prod(shape, dtype=int))).reshape(shape)
        for pattern in [
            "a b c d e f -> a b c d e f",
            "b a c d e f -> a b d e f c",
            "a b c d e f -> f e d c b a",
            "a b c d e f -> (f e) d (c b a)",
            "a b c d e f -> (f e d c b a)",
        ]:
            result = rearrange(x, pattern)
            self.assertEqual(len(np.setdiff1d(x, result)), 0)
            self.assertIs(result.dtype, x.dtype)

        result = rearrange(x, "a b c d e f -> a (b) (c d e) f")
        self.assertEqual(x.flatten(), result.flatten())

        result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11")
        self.assertEqual(x, result)

        result1 = rearrange(x, "a b c d e f -> f e d c b a")
        result2 = rearrange(x, "f e d c b a -> a b c d e f")
        self.assertEqual(result1, result2)

        result = rearrange(
            rearrange(x, "a b c d e f -> (f d) c (e b) a"),
            "(f d) c (e b) a -> a b c d e f",
            b=2,
            d=5,
        )
        self.assertEqual(x, result)

        sizes = dict(zip("abcdef", shape))
        temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes)
        result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes)
        self.assertEqual(x, result)

        x2 = torch.arange(2 * 3 * 4).reshape([2, 3, 4])
        result = rearrange(x2, "a b c -> b c a")
        self.assertEqual(x2[1, 2, 3], result[2, 3, 1])
        self.assertEqual(x2[0, 1, 2], result[1, 2, 0])
//  [python]
def test_rearrange_consistency(self) -> None:
        shape = [1, 2, 3, 5, 7, 11]
        x = torch.arange(int(np.prod(shape, dtype=int))).reshape(shape)
        for pattern in [
            "a b c d e f -> a b c d e f",
            "b a c d e f -> a b d e f c",
            "a b c d e f -> f e d c b a",
            "a b c d e f -> (f e) d (c b a)",
            "a b c d e f -> (f e d c b a)",
        ]:
            result = rearrange(x, pattern)
            self.assertEqual(len(np.setdiff1d(x, result)), 0)
            self.assertIs(result.dtype, x.dtype)

        result = rearrange(x, "a b c d e f -> a (b) (c d e) f")
        torch.testing.assert_close(x.flatten(), result.flatten())

        result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11")
        torch.testing.assert_close(x, result)

        result1 = rearrange(x, "a b c d e f -> f e d c b a")
        result2 = rearrange(x, "f e d c b a -> a b c d e f")
        torch.testing.assert_close(result1, result2)

        result = rearrange(
            rearrange(x, "a b c d e f -> (f d) c (e b) a"),
            "(f d) c (e b) a -> a b c d e f",
            b=2,
            d=5,
        )
        torch.testing.assert_close(x, result)

        sizes = dict(zip("abcdef", shape))
        temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes)
        result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes)
        torch.testing.assert_close(x, result)

        x2 = torch.arange(2 * 3 * 4).reshape([2, 3, 4])
        result = rearrange(x2, "a b c -> b c a")
        self.assertEqual(x2[1, 2, 3], result[2, 3, 1])
        self.assertEqual(x2[0, 1, 2], result[1, 2, 0])
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BERT model to flash_attn and back.
    """

    state_dict = BertForPreTrainingHF.from_pretrained(model_name).state_dict()
    config = BertConfig.from_pretrained(model_name)

    flash_state_dict = remap_bert_state_dict(state_dict, config, add_pooling_layer=True)
    recovered_state_dict = inv_remap_state_dict(flash_state_dict, config)

    assert (set(state_dict.keys()) - set(['cls.seq_relationship.bias', 'cls.seq_relationship.weight'])) == set(
        recovered_state_dict.keys()
    )

    for k in state_dict.keys():
        if k in ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']:
            continue
        assert state_dict[k].shape == recovered_state_dict[k].shape
        t.assert_close(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BERT model to flash_attn and back.
    """

    state_dict = BertForPreTrainingHF.from_pretrained(model_name).state_dict()
    config = BertConfig.from_pretrained(model_name)

    flash_state_dict = remap_bert_state_dict(state_dict, config, add_pooling_layer=True)
    recovered_state_dict = inv_remap_state_dict(flash_state_dict, config)

    assert (set(state_dict.keys()) - set(['cls.seq_relationship.bias', 'cls.seq_relationship.weight'])) == set(
        recovered_state_dict.keys()
    )

    for k in state_dict.keys():
        if k in ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']:
            continue
        assert state_dict[k].shape == recovered_state_dict[k].shape
        torch.testing.assert_close(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def test_qkv_split():
    # MHA
    config = Config(n_embd=4, n_head=4)
    qkv = torch.tensor(
        [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [8, 9, 10, 11],
            [12, 13, 14, 15],
            [16, 17, 18, 19],
            [20, 21, 22, 23],
            [24, 25, 26, 27],
            [28, 29, 30, 31],
            [32, 33, 34, 35],
            [36, 37, 38, 39],
            [40, 41, 42, 43],
            [44, 45, 46, 47],
        ]
    )
    q, k, v = qkv_split(qkv, config)
    t.assert_close(q, torch.tensor([[0, 1, 2, 3], [12, 13, 14, 15], [24, 25, 26, 27], [36, 37, 38, 39]]))
    t.assert_close(k, torch.tensor([[4, 5, 6, 7], [16, 17, 18, 19], [28, 29, 30, 31], [40, 41, 42, 43]]))
    t.assert_close(v, torch.tensor([[8, 9, 10, 11], [20, 21, 22, 23], [32, 33, 34, 35], [44, 45, 46, 47]]))

    # GQA
    config = Config(n_embd=4, n_head=4, n_query_groups=2)
    qkv = torch.tensor(
        [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [8, 9, 10, 11],
            [12, 13, 14, 15],
            [16, 17, 18, 19],
            [20, 21, 22, 23],
            [24, 25, 26, 27],
            [28, 29, 30, 31],
        ]
    )
    q, k, v = qkv_split(qkv, config)
    t.assert_close(q, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7], [16, 17, 18, 19], [20, 21, 22, 23]]))
    t.assert_close(k, torch.tensor([[8, 9, 10, 11], [24, 25, 26, 27]]))
    t.assert_close(v, torch.tensor([[12, 13, 14, 15], [28, 29, 30, 31]]))

    # MQA
    config = Config(n_embd=4, n_head=4, n_query_groups=1)
    qkv = torch.tensor(
        [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]
    )
    q, k, v = qkv_split(qkv, config)
    t.assert_close(q, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]))
    t.assert_close(k, torch.tensor([[16, 17, 18, 19]]))
    t.assert_close(v, torch.tensor([[20, 21, 22, 23]]))
//  [python]
def test_qkv_split():
    # MHA
    config = Config(n_embd=4, n_head=4)
    qkv = torch.tensor(
        [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [8, 9, 10, 11],
            [12, 13, 14, 15],
            [16, 17, 18, 19],
            [20, 21, 22, 23],
            [24, 25, 26, 27],
            [28, 29, 30, 31],
            [32, 33, 34, 35],
            [36, 37, 38, 39],
            [40, 41, 42, 43],
            [44, 45, 46, 47],
        ]
    )
    q, k, v = qkv_split(qkv, config)
    torch.testing.assert_close(q, torch.tensor([[0, 1, 2, 3], [12, 13, 14, 15], [24, 25, 26, 27], [36, 37, 38, 39]]))
    torch.testing.assert_close(k, torch.tensor([[4, 5, 6, 7], [16, 17, 18, 19], [28, 29, 30, 31], [40, 41, 42, 43]]))
    torch.testing.assert_close(v, torch.tensor([[8, 9, 10, 11], [20, 21, 22, 23], [32, 33, 34, 35], [44, 45, 46, 47]]))

    # GQA
    config = Config(n_embd=4, n_head=4, n_query_groups=2)
    qkv = torch.tensor(
        [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [8, 9, 10, 11],
            [12, 13, 14, 15],
            [16, 17, 18, 19],
            [20, 21, 22, 23],
            [24, 25, 26, 27],
            [28, 29, 30, 31],
        ]
    )
    q, k, v = qkv_split(qkv, config)
    torch.testing.assert_close(q, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7], [16, 17, 18, 19], [20, 21, 22, 23]]))
    torch.testing.assert_close(k, torch.tensor([[8, 9, 10, 11], [24, 25, 26, 27]]))
    torch.testing.assert_close(v, torch.tensor([[12, 13, 14, 15], [28, 29, 30, 31]]))

    # MQA
    config = Config(n_embd=4, n_head=4, n_query_groups=1)
    qkv = torch.tensor(
        [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]
    )
    q, k, v = qkv_split(qkv, config)
    torch.testing.assert_close(q, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]))
    torch.testing.assert_close(k, torch.tensor([[16, 17, 18, 19]]))
    torch.testing.assert_close(v, torch.tensor([[20, 21, 22, 23]]))
//  [python]
def test_actorcritic(device):
    common_module = SafeModule(
        module=nn.Linear(3, 4), in_keys=["obs"], out_keys=["hidden"], spec=None
    ).to(device)
    module = SafeModule(nn.Linear(4, 5), in_keys=["hidden"], out_keys=["param"])
    policy_operator = ProbabilisticActor(
        module=module, in_keys=["param"], spec=None, return_log_prob=True
    ).to(device)
    value_operator = ValueOperator(nn.Linear(4, 1), in_keys=["hidden"]).to(device)
    op = ActorValueOperator(
        common_operator=common_module,
        policy_operator=policy_operator,
        value_operator=value_operator,
    ).to(device)
    td = TensorDict(
        source={"obs": torch.randn(4, 3)},
        batch_size=[
            4,
        ],
    ).to(device)
    td_total = op(td.clone())
    policy_op = op.get_policy_operator()
    td_policy = policy_op(td.clone())
    value_op = op.get_value_operator()
    td_value = value_op(td)
    torch.testing.assert_allclose(td_total.get("action"), td_policy.get("action"))
    torch.testing.assert_allclose(
        td_total.get("sample_log_prob"), td_policy.get("sample_log_prob")
    )
    torch.testing.assert_allclose(td_total.get("state_value"), td_value.get("state_value"))

    value_params = set(
        list(op.get_value_operator().parameters()) + list(op.module[0].parameters())
    )
    value_params2 = set(value_op.parameters())
    assert len(value_params.difference(value_params2)) == 0 and len(
        value_params.intersection(value_params2)
    ) == len(value_params)

    policy_params = set(
        list(op.get_policy_operator().parameters()) + list(op.module[0].parameters())
    )
    policy_params2 = set(policy_op.parameters())
    assert len(policy_params.difference(policy_params2)) == 0 and len(
        policy_params.intersection(policy_params2)
    ) == len(policy_params)
//  [python]
def test_actorcritic(device):
    common_module = SafeModule(
        module=nn.Linear(3, 4), in_keys=["obs"], out_keys=["hidden"], spec=None
    ).to(device)
    module = SafeModule(nn.Linear(4, 5), in_keys=["hidden"], out_keys=["param"])
    policy_operator = ProbabilisticActor(
        module=module, in_keys=["param"], spec=None, return_log_prob=True
    ).to(device)
    value_operator = ValueOperator(nn.Linear(4, 1), in_keys=["hidden"]).to(device)
    op = ActorValueOperator(
        common_operator=common_module,
        policy_operator=policy_operator,
        value_operator=value_operator,
    ).to(device)
    td = TensorDict(
        source={"obs": torch.randn(4, 3)},
        batch_size=[
            4,
        ],
    ).to(device)
    td_total = op(td.clone())
    policy_op = op.get_policy_operator()
    td_policy = policy_op(td.clone())
    value_op = op.get_value_operator()
    td_value = value_op(td)
    torch.testing.assert_close(td_total.get("action"), td_policy.get("action"))
    torch.testing.assert_close(
        td_total.get("sample_log_prob"), td_policy.get("sample_log_prob")
    )
    torch.testing.assert_close(td_total.get("state_value"), td_value.get("state_value"))

    value_params = set(
        list(op.get_value_operator().parameters()) + list(op.module[0].parameters())
    )
    value_params2 = set(value_op.parameters())
    assert len(value_params.difference(value_params2)) == 0 and len(
        value_params.intersection(value_params2)
    ) == len(value_params)

    policy_params = set(
        list(op.get_policy_operator().parameters()) + list(op.module[0].parameters())
    )
    policy_params2 = set(policy_op.parameters())
    assert len(policy_params.difference(policy_params2)) == 0 and len(
        policy_params.intersection(policy_params2)
    ) == len(policy_params)
//  [python]
def test_batched_intervention_2d_graph(graph, intervention_variable, intervention_value, two_variable_dict):
    rff_features = torch.rand((10, 3))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(two_variable_dict, rff_features, coef_matrix, graph)
    variable_names = set(two_variable_dict.keys())
    sampled_variables = variable_names - {intervention_variable}
    assert len(sampled_variables) == 1
    sampled_variable = sampled_variables.pop()
    intervention_value = torch.tensor(intervention_value)
    batched_do_sem = sem.do(
        TensorDict(
            {intervention_variable: intervention_value},
            batch_size=[
                3,
            ],
        )
    )
    noise = batched_do_sem.sample_noise((10,))
    do_sample = batched_do_sem.noise_to_sample(noise)

    non_batch_sample_list = []
    for i, intervention in enumerate(intervention_value):
        non_batch_sample_list.append(
            sem.do(TensorDict({intervention_variable: intervention}, batch_size=tuple()))
            .noise_to_sample(noise[:, i, None])
            .squeeze(1)
        )
    non_batch_sample = torch.stack(non_batch_sample_list, dim=1)
    torch.testing.assert_allclose(do_sample[sampled_variable], non_batch_sample[sampled_variable], atol=1e-5, rtol=1e-4)
//  [python]
def test_batched_intervention_2d_graph(graph, intervention_variable, intervention_value, two_variable_dict):
    rff_features = torch.rand((10, 3))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(two_variable_dict, rff_features, coef_matrix, graph)
    variable_names = set(two_variable_dict.keys())
    sampled_variables = variable_names - {intervention_variable}
    assert len(sampled_variables) == 1
    sampled_variable = sampled_variables.pop()
    intervention_value = torch.tensor(intervention_value)
    batched_do_sem = sem.do(
        TensorDict(
            {intervention_variable: intervention_value},
            batch_size=[
                3,
            ],
        )
    )
    noise = batched_do_sem.sample_noise((10,))
    do_sample = batched_do_sem.noise_to_sample(noise)

    non_batch_sample_list = []
    for i, intervention in enumerate(intervention_value):
        non_batch_sample_list.append(
            sem.do(TensorDict({intervention_variable: intervention}, batch_size=tuple()))
            .noise_to_sample(noise[:, i, None])
            .squeeze(1)
        )
    non_batch_sample = torch.stack(non_batch_sample_list, dim=1)
    torch.testing.assert_close(do_sample[sampled_variable], non_batch_sample[sampled_variable], atol=1e-5, rtol=1e-4)
//  [python]
def test_batched_intervention_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graph)
    intervention_value = torch.tensor([[1.42, 0.42], [0.42, 1.42], [0.42, 0.42]])
    batched_do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value},
            batch_size=[
                3,
            ],
        )
    )
    noise = batched_do_sem.sample_noise((10,))
    do_sample = batched_do_sem.noise_to_sample(noise)
    inferred_noise = batched_do_sem.sample_to_noise(do_sample)
    t.assert_close(noise["x1"], inferred_noise["x1"], atol=1e-5, rtol=1e-4)
    t.assert_close(noise["x3"], inferred_noise["x3"], atol=1e-5, rtol=1e-4)

    non_batch_sample_list = []
    for i, intervention in enumerate(intervention_value):
        non_batch_sample_list.append(
            sem.do(TensorDict({"x2": intervention}, batch_size=tuple())).noise_to_sample(noise[:, i, None]).squeeze(1)
        )

    non_batch_sample = torch.stack(non_batch_sample_list, dim=1)

    t.assert_close(do_sample["x1"], non_batch_sample["x1"], atol=1e-5, rtol=1e-4)
    t.assert_close(do_sample["x3"], non_batch_sample["x3"], atol=1e-5, rtol=1e-4)
//  [python]
def test_batched_intervention_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graph)
    intervention_value = torch.tensor([[1.42, 0.42], [0.42, 1.42], [0.42, 0.42]])
    batched_do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value},
            batch_size=[
                3,
            ],
        )
    )
    noise = batched_do_sem.sample_noise((10,))
    do_sample = batched_do_sem.noise_to_sample(noise)
    inferred_noise = batched_do_sem.sample_to_noise(do_sample)
    torch.testing.assert_close(noise["x1"], inferred_noise["x1"], atol=1e-5, rtol=1e-4)
    torch.testing.assert_close(noise["x3"], inferred_noise["x3"], atol=1e-5, rtol=1e-4)

    non_batch_sample_list = []
    for i, intervention in enumerate(intervention_value):
        non_batch_sample_list.append(
            sem.do(TensorDict({"x2": intervention}, batch_size=tuple())).noise_to_sample(noise[:, i, None]).squeeze(1)
        )

    non_batch_sample = torch.stack(non_batch_sample_list, dim=1)

    torch.testing.assert_close(do_sample["x1"], non_batch_sample["x1"], atol=1e-5, rtol=1e-4)
    torch.testing.assert_close(do_sample["x3"], non_batch_sample["x3"], atol=1e-5, rtol=1e-4)
//  [python]
def test_do_linear_sem(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    sem = create_lingauss_sem(two_variable_dict, coef_matrix, graph)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x2": intervention_value}, batch_size=tuple()))
    array = torch.linspace(-2, 2, 100).unsqueeze(-1)
    log_probs = do_sem.log_prob(TensorDict({"x1": array}, batch_size=[100]))
    if graph[1, 0] == 1.0:
        expected_mean = torch.einsum("i,ij->j", intervention_value, coef_matrix[1:, :1])
    else:
        expected_mean = torch.tensor([0.0])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(array)
    t.assert_close(log_probs, expected_log_probs)
    noise = sem.sample_noise((10,))
    do_sample = do_sem.noise_to_sample(noise)
    sample = sem.noise_to_sample(noise)
    if graph[1, 0] == 1.0:
        t.assert_close(do_sample["x1"], expected_mean + noise["x1"])
    else:
        t.assert_close(do_sample["x1"], sample["x1"])

    # Test multi-dim interventions
    do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value.expand(3, 2)},
            batch_size=[
                3,
            ],
        )
    )
    noise = sem.sample_noise((10, 3, 1))
    do_sample = do_sem.noise_to_sample(noise)
    sample = sem.noise_to_sample(noise)
    if graph[1, 0] == 1.0:
        t.assert_close(do_sample["x1"], expected_mean + noise["x1"])
    else:
        t.assert_close(do_sample["x1"], sample["x1"])
//  [python]
def test_do_linear_sem(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    sem = create_lingauss_sem(two_variable_dict, coef_matrix, graph)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x2": intervention_value}, batch_size=tuple()))
    array = torch.linspace(-2, 2, 100).unsqueeze(-1)
    log_probs = do_sem.log_prob(TensorDict({"x1": array}, batch_size=[100]))
    if graph[1, 0] == 1.0:
        expected_mean = torch.einsum("i,ij->j", intervention_value, coef_matrix[1:, :1])
    else:
        expected_mean = torch.tensor([0.0])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(array)
    torch.testing.assert_close(log_probs, expected_log_probs)
    noise = sem.sample_noise((10,))
    do_sample = do_sem.noise_to_sample(noise)
    sample = sem.noise_to_sample(noise)
    if graph[1, 0] == 1.0:
        torch.testing.assert_close(do_sample["x1"], expected_mean + noise["x1"])
    else:
        torch.testing.assert_close(do_sample["x1"], sample["x1"])

    # Test multi-dim interventions
    do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value.expand(3, 2)},
            batch_size=[
                3,
            ],
        )
    )
    noise = sem.sample_noise((10, 3, 1))
    do_sample = do_sem.noise_to_sample(noise)
    sample = sem.noise_to_sample(noise)
    if graph[1, 0] == 1.0:
        torch.testing.assert_close(do_sample["x1"], expected_mean + noise["x1"])
    else:
        torch.testing.assert_close(do_sample["x1"], sample["x1"])
//  [python]
def test_lmhead_actorvalueoperator(device):
    from transformers import AutoModelForCausalLM, GPT2Config

    config = GPT2Config(return_dict=False)
    base_model = AutoModelForCausalLM.from_config(config).eval()
    aco = LMHeadActorValueOperator(base_model).to(device)

    # check common
    assert aco.module[0][0].module is base_model.transformer
    assert aco.module[0][1].in_keys == ["x"]
    assert aco.module[0][1].out_keys == ["x"]

    # check actor
    assert aco.module[1].in_keys == ["x"]
    assert aco.module[1].out_keys == ["logits", "action", "sample_log_prob"]
    assert aco.module[1][0].module is base_model.lm_head

    # check critic
    assert aco.module[2].in_keys == ["x"]
    assert aco.module[2].out_keys == ["state_value"]
    assert isinstance(aco.module[2].module, nn.Linear)
    assert aco.module[2].module.in_features == base_model.transformer.embed_dim
    assert aco.module[2].module.out_features == 1

    td = TensorDict(
        source={
            "input_ids": torch.randint(50257, (4, 3)),
            "attention_mask": torch.ones((4, 3)),
        },
        batch_size=[
            4,
        ],
        device=device,
    )
    td_total = aco(td.clone())
    policy_op = aco.get_policy_operator()
    td_policy = policy_op(td.clone())
    value_op = aco.get_value_operator()
    td_value = value_op(td)
    torch.testing.assert_allclose(td_total.get("action"), td_policy.get("action"))
    torch.testing.assert_allclose(
        td_total.get("sample_log_prob"), td_policy.get("sample_log_prob")
    )
    torch.testing.assert_allclose(td_total.get("state_value"), td_value.get("state_value"))

    value_params = set(
        list(aco.get_value_operator().parameters()) + list(aco.module[0].parameters())
    )
    value_params2 = set(value_op.parameters())
    assert len(value_params.difference(value_params2)) == 0 and len(
        value_params.intersection(value_params2)
    ) == len(value_params)

    policy_params = set(
        list(aco.get_policy_operator().parameters()) + list(aco.module[0].parameters())
    )
    policy_params2 = set(policy_op.parameters())
    assert len(policy_params.difference(policy_params2)) == 0 and len(
        policy_params.intersection(policy_params2)
    ) == len(policy_params)