jiahuimbzuai/CoQuIR huggingface dataset:
//  [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_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(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)
//  [python]
def test_intervention_batched_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    graphs = torch.stack([graph, graph], dim=0)
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graphs)

    regular_noise = sem.sample_noise(torch.Size([10]))
    samples = sem.noise_to_sample(regular_noise)
    inferred_noise = sem.sample_to_noise(samples)

    t.assert_close(regular_noise, inferred_noise, atol=1e-5, rtol=1e-4)

    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value},
            batch_size=torch.Size([]),
        )
    )
    noise = do_sem.sample_noise(torch.Size([10]))
    do_sample = do_sem.noise_to_sample(noise)
    inferred_noise = 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)
//  [python]
def test_intervention_batched_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    graphs = torch.stack([graph, graph], dim=0)
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graphs)

    regular_noise = sem.sample_noise(torch.Size([10]))
    samples = sem.noise_to_sample(regular_noise)
    inferred_noise = sem.sample_to_noise(samples)

    torch.testing.assert_close(regular_noise, inferred_noise, atol=1e-5, rtol=1e-4)

    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(
        TensorDict(
            {"x2": intervention_value},
            batch_size=torch.Size([]),
        )
    )
    noise = do_sem.sample_noise(torch.Size([10]))
    do_sample = do_sem.noise_to_sample(noise)
    inferred_noise = 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)
//  [python]
def test_linear_sem_3d_graph_do_1_node(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x1": intervention_value}, batch_size=tuple()))
    test_val = torch.rand(100, 3)
    log_probs = do_sem.log_prob(TensorDict({"x2": test_val[:, 0:2], "x3": test_val[:, 2:]}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", intervention_value, coef_matrix[:2, 2:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    t.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_1_node(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x1": intervention_value}, batch_size=tuple()))
    test_val = torch.rand(100, 3)
    log_probs = do_sem.log_prob(TensorDict({"x2": test_val[:, 0:2], "x3": test_val[:, 2:]}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", intervention_value, coef_matrix[:2, 2:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    torch.testing.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_2_nodes(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.triu(torch.ones(3, 3), diagonal=1)
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    int_data = TensorDict({"x1": torch.tensor([1.42, 0.42]), "x2": torch.tensor([-0.42, 0.402])}, batch_size=tuple())
    do_sem = sem.do(int_data)
    test_val = torch.linspace(-2, 2, 100).unsqueeze(-1)
    log_probs = do_sem.log_prob(TensorDict({"x3": test_val}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", torch.cat((int_data["x1"], int_data["x2"])), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    t.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_2_nodes(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.triu(torch.ones(3, 3), diagonal=1)
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    int_data = TensorDict({"x1": torch.tensor([1.42, 0.42]), "x2": torch.tensor([-0.42, 0.402])}, batch_size=tuple())
    do_sem = sem.do(int_data)
    test_val = torch.linspace(-2, 2, 100).unsqueeze(-1)
    log_probs = do_sem.log_prob(TensorDict({"x3": test_val}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", torch.cat((int_data["x1"], int_data["x2"])), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    torch.testing.assert_close(log_probs, expected_log_probs)
//  [python]
def test_process_transform_is_not_none(self, data):
        dm = DataModel(data=data, transform=to_tensor)
        t.assert_allclose(dm.process(data[0]), torch.tensor([1, 0, 0, 1]))
        t.assert_allclose(dm.process(data[1]), torch.tensor([0, 1, 1, 0]))
//  [python]
def test_process_transform_is_not_none(self, data):
        dm = DataModel(data=data, transform=to_tensor)
        torch.testing.assert_close(dm.process(data[0]), torch.tensor([1, 0, 0, 1]))
        torch.testing.assert_close(dm.process(data[1]), torch.tensor([0, 1, 1, 0]))
//  [python]
def test_batched_intervention_batched_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    graphs = torch.stack([graph, graph], dim=0)
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graphs)
    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)
    assert do_sample.batch_size == torch.Size([10, 3, 2])  # 10 samples, 3 interventions, 2 graphs
    inferred_noise = batched_do_sem.sample_to_noise(do_sample)
    torch.allclose(noise["x1"], inferred_noise["x1"], atol=1e-5, rtol=1e-4)
    torch.allclose(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.allclose(do_sample["x1"], non_batch_sample["x1"], atol=1e-5, rtol=1e-4)
    torch.allclose(do_sample["x3"], non_batch_sample["x3"], atol=1e-5, rtol=1e-4)
//  [python]
def test_batched_intervention_batched_3d_graph(three_variable_dict):
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    graphs = torch.stack([graph, graph], dim=0)
    rff_features = torch.rand((10, 5))
    coef_matrix = torch.rand((10,))
    sem = create_rffgauss_sem(three_variable_dict, rff_features, coef_matrix, graphs)
    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)
    assert do_sample.batch_size == torch.Size([10, 3, 2])  # 10 samples, 3 interventions, 2 graphs
    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_linear_sem_3d_graph_condition_1_node(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    condition_value = torch.tensor([1.42, 0.42])
    condition_sem = sem.condition(TensorDict({"x1": condition_value}, batch_size=tuple()))
    test_val = torch.rand(100, 3)
    log_probs = condition_sem.log_prob(TensorDict({"x2": test_val[:, 0:2], "x3": test_val[:, 2:]}, batch_size=[100]))

    with pytest.raises(NotImplementedError, match=r"on x2"):
        sem.condition(TensorDict({"x2": condition_value}, batch_size=tuple()))

    condition_sem = sem.condition(TensorDict({"x1": condition_value, "x2": condition_value}, batch_size=tuple()))
    test_val = torch.rand(100, 1)
    log_probs = condition_sem.log_prob(TensorDict({"x3": test_val}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", torch.cat([condition_value, condition_value]), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    t.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_condition_1_node(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = graph[1, 2] = 1
    sem = create_lingauss_sem(three_variable_dict, coef_matrix, graph)
    condition_value = torch.tensor([1.42, 0.42])
    condition_sem = sem.condition(TensorDict({"x1": condition_value}, batch_size=tuple()))
    test_val = torch.rand(100, 3)
    log_probs = condition_sem.log_prob(TensorDict({"x2": test_val[:, 0:2], "x3": test_val[:, 2:]}, batch_size=[100]))

    with pytest.raises(NotImplementedError, match=r"on x2"):
        sem.condition(TensorDict({"x2": condition_value}, batch_size=tuple()))

    condition_sem = sem.condition(TensorDict({"x1": condition_value, "x2": condition_value}, batch_size=tuple()))
    test_val = torch.rand(100, 1)
    log_probs = condition_sem.log_prob(TensorDict({"x3": test_val}, batch_size=[100]))
    expected_mean = torch.einsum("i,ij->j", torch.cat([condition_value, condition_value]), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Normal(expected_mean, 1.0), 1).log_prob(test_val)
    torch.testing.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_1_node_bernoulli(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = 1
    func = LinearFunctionalRelationships(three_variable_dict, coef_matrix)
    noise_modules = create_noise_modules(
        shapes=three_variable_dict,
        types=dict.fromkeys(three_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x1": intervention_value}, batch_size=tuple()))
    array = torch.bernoulli(0.5 * torch.ones(100, 3))
    log_probs = do_sem.log_prob(TensorDict({"x2": array[:, :2], "x3": array[:, 2:]}, batch_size=array.shape[:-1]))
    expected_logits = torch.einsum("i,ij->j", intervention_value, coef_matrix[:2, 2:])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 1).log_prob(array)
    torch.testing.assert_allclose(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_1_node_bernoulli(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.zeros(3, 3)
    graph[0, 1] = graph[0, 2] = 1
    func = LinearFunctionalRelationships(three_variable_dict, coef_matrix)
    noise_modules = create_noise_modules(
        shapes=three_variable_dict,
        types=dict.fromkeys(three_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x1": intervention_value}, batch_size=tuple()))
    array = torch.bernoulli(0.5 * torch.ones(100, 3))
    log_probs = do_sem.log_prob(TensorDict({"x2": array[:, :2], "x3": array[:, 2:]}, batch_size=array.shape[:-1]))
    expected_logits = torch.einsum("i,ij->j", intervention_value, coef_matrix[:2, 2:])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 1).log_prob(array)
    torch.testing.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_2_nodes_bernoulli(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.triu(torch.ones(3, 3), diagonal=1)
    func = LinearFunctionalRelationships(three_variable_dict, coef_matrix)
    noise_modules = create_noise_modules(
        shapes=three_variable_dict,
        types=dict.fromkeys(three_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    int_data = TensorDict({"x1": torch.tensor([1.42, 0.42]), "x2": torch.tensor([-0.42, 0.402])}, batch_size=tuple())
    do_sem = sem.do(int_data)
    array = torch.bernoulli(0.5 * torch.ones(100, 1))
    log_probs = do_sem.log_prob(TensorDict({"x3": array}, batch_size=array.shape[:-1]))
    expected_logits = torch.einsum("i,ij->j", torch.cat((int_data["x1"], int_data["x2"])), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 1).log_prob(array)
    t.assert_close(log_probs, expected_log_probs)
//  [python]
def test_linear_sem_3d_graph_do_2_nodes_bernoulli(three_variable_dict):
    coef_matrix = torch.rand((5, 5))
    graph = torch.triu(torch.ones(3, 3), diagonal=1)
    func = LinearFunctionalRelationships(three_variable_dict, coef_matrix)
    noise_modules = create_noise_modules(
        shapes=three_variable_dict,
        types=dict.fromkeys(three_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    int_data = TensorDict({"x1": torch.tensor([1.42, 0.42]), "x2": torch.tensor([-0.42, 0.402])}, batch_size=tuple())
    do_sem = sem.do(int_data)
    array = torch.bernoulli(0.5 * torch.ones(100, 1))
    log_probs = do_sem.log_prob(TensorDict({"x3": array}, batch_size=array.shape[:-1]))
    expected_logits = torch.einsum("i,ij->j", torch.cat((int_data["x1"], int_data["x2"])), coef_matrix[:4, 4:])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 1).log_prob(array)
    torch.testing.assert_close(log_probs, expected_log_probs)
//  [python]
def test_do_linear_sem_bernoulli(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    func = LinearFunctionalRelationships(two_variable_dict, initial_linear_coefficient_matrix=coef_matrix)
    noise_modules = create_noise_modules(
        shapes=two_variable_dict,
        types=dict.fromkeys(two_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x2": intervention_value}, batch_size=tuple()))
    array = torch.bernoulli(0.5 * torch.ones(100, 1))
    log_probs = do_sem.log_prob(TensorDict({"x1": array}, batch_size=array.shape[:-1]))
    if graph[1, 0] == 1.0:
        expected_logits = torch.einsum("i,ij->j", intervention_value, coef_matrix[1:, :1])
    else:
        expected_logits = torch.tensor([0.0])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 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_logits + noise["x1"]) > 0).float())
    else:
        t.assert_close(do_sample["x1"], sample["x1"])
//  [python]
def test_do_linear_sem_bernoulli(graph, two_variable_dict):
    coef_matrix = torch.rand((3, 3))
    func = LinearFunctionalRelationships(two_variable_dict, initial_linear_coefficient_matrix=coef_matrix)
    noise_modules = create_noise_modules(
        shapes=two_variable_dict,
        types=dict.fromkeys(two_variable_dict, VariableTypeEnum.BINARY),
        continuous_noise_dist=ContinuousNoiseDist.GAUSSIAN,
    )
    noise_dist = JointNoiseModule(noise_modules)
    sem = DistributionParametersSEM(graph=graph, noise_dist=noise_dist, func=func)
    intervention_value = torch.tensor([1.42, 0.42])
    do_sem = sem.do(TensorDict({"x2": intervention_value}, batch_size=tuple()))
    array = torch.bernoulli(0.5 * torch.ones(100, 1))
    log_probs = do_sem.log_prob(TensorDict({"x1": array}, batch_size=array.shape[:-1]))
    if graph[1, 0] == 1.0:
        expected_logits = torch.einsum("i,ij->j", intervention_value, coef_matrix[1:, :1])
    else:
        expected_logits = torch.tensor([0.0])
    expected_log_probs = td.Independent(td.Bernoulli(logits=expected_logits), 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_logits + noise["x1"]) > 0).float())
    else:
        torch.testing.assert_close(do_sample["x1"], sample["x1"])
//  [python]
def unscale_l2norm(self, sizea, sizeb, applier, repeat_tensors, in_type, per_tensor):
        self.overflow_buf.zero_()
        a = torch.full([sizea], self.val, dtype=torch.float32, device='cuda')
        b = torch.full([sizeb], self.val, dtype=torch.float32, device='cuda')

        in_list = []
        for i in range(repeat_tensors):
            in_list += [a.clone().to(in_type), b.clone().to(in_type)]

        if per_tensor:
            norm, norm_per_tensor = applier(multi_tensor_unscale_l2norm, self.overflow_buf, [in_list], self.inv_scale_cuda, True)
            normab = torch.cat(((a * self.inv_scale).norm().view(1), (b * self.inv_scale).norm().view(1)))
            norm_per_tensor = norm_per_tensor.view(-1, 2)
        else:
            norm, _ = applier(multi_tensor_unscale_l2norm, self.overflow_buf, [in_list], self.inv_scale_cuda, True)

        reference = torch.full([(sizea + sizeb)*repeat_tensors], self.val * self.inv_scale, dtype=torch.float32, device='cuda').norm()

        self.assertEqual(norm, reference)
        if per_tensor:
            self.assertEqual(norm_per_tensor, normab)
        self.assertTrue(self.overflow_buf.item() == 0)
//  [python]
def unscale_l2norm(self, sizea, sizeb, applier, repeat_tensors, in_type, per_tensor):
        self.overflow_buf.zero_()
        a = torch.full([sizea], self.val, dtype=torch.float32, device='cuda')
        b = torch.full([sizeb], self.val, dtype=torch.float32, device='cuda')

        in_list = []
        for i in range(repeat_tensors):
            in_list += [a.clone().to(in_type), b.clone().to(in_type)]

        if per_tensor:
            norm, norm_per_tensor = applier(multi_tensor_unscale_l2norm, self.overflow_buf, [in_list], self.inv_scale_cuda, True)
            normab = torch.cat(((a * self.inv_scale).norm().view(1), (b * self.inv_scale).norm().view(1)))
            norm_per_tensor = norm_per_tensor.view(-1, 2)
        else:
            norm, _ = applier(multi_tensor_unscale_l2norm, self.overflow_buf, [in_list], self.inv_scale_cuda, True)

        reference = torch.full([(sizea + sizeb)*repeat_tensors], self.val * self.inv_scale, dtype=torch.float32, device='cuda').norm()

        torch.testing.assert_close(norm, reference)
        if per_tensor:
            torch.testing.assert_close(norm_per_tensor, normab)
        self.assertTrue(self.overflow_buf.item() == 0)
//  [python]
def test_auc_input_dim_1(self) -> None:
        x = torch.rand(3)
        y = torch.rand(1, 3)

        sklearn_x, sklearn_x_idx = x.unsqueeze(0).sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y, reorder=True)
        torch.testing.assert_allclose(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_input_dim_1(self) -> None:
        x = torch.rand(3)
        y = torch.rand(1, 3)

        sklearn_x, sklearn_x_idx = x.unsqueeze(0).sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y, reorder=True)
        torch.testing.assert_close(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_base(self) -> None:
        x = torch.tensor(
            [
                [0.04, 0.06, 0.8],
                [0.1, 0.2, 0.7],
                [0.04, 0.06, 0.08],
                [0.10, 0.14, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])
        sklearn_result = torch.tensor(
            [sklearn_auc(a, b) for b, a in zip(y, x)], dtype=torch.float32
        )
        torch_auc_result = auc(x, y)
        torch.testing.assert_allclose(
            torch_auc_result,
            sklearn_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_base(self) -> None:
        x = torch.tensor(
            [
                [0.04, 0.06, 0.8],
                [0.1, 0.2, 0.7],
                [0.04, 0.06, 0.08],
                [0.10, 0.14, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])
        sklearn_result = torch.tensor(
            [sklearn_auc(a, b) for b, a in zip(y, x)], dtype=torch.float32
        )
        torch_auc_result = auc(x, y)
        torch.testing.assert_close(
            torch_auc_result,
            sklearn_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_reordered_input(self) -> None:
        x = torch.tensor(
            [
                [0.05, 0.02, 1.8],
                [0.6, 0.3, 0.7],
                [0.02, 0.01, 0.08],
                [0.15, 0.11, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])

        sklearn_x, sklearn_x_idx = x.sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y, reorder=True)
        torch.testing.assert_allclose(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_reordered_input(self) -> None:
        x = torch.tensor(
            [
                [0.05, 0.02, 1.8],
                [0.6, 0.3, 0.7],
                [0.02, 0.01, 0.08],
                [0.15, 0.11, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])

        sklearn_x, sklearn_x_idx = x.sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y, reorder=True)
        torch.testing.assert_close(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_no_reorder(self) -> None:
        x = torch.tensor(
            [
                [0.02, 0.05, 1.8],
                [0.3, 0.4, 0.7],
                [0.03, 0.05, 0.08],
                [0.10, 0.11, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])

        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(y, x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y)
        torch.testing.assert_allclose(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_no_reorder(self) -> None:
        x = torch.tensor(
            [
                [0.02, 0.05, 1.8],
                [0.3, 0.4, 0.7],
                [0.03, 0.05, 0.08],
                [0.10, 0.11, 0.18],
            ]
        )
        y = torch.tensor([[2, 0, 1], [0, 1, 0], [1, 1, 1], [2, 0, 1]])

        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(y, x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y)
        torch.testing.assert_close(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_mean(preds, expected_mean):
    x = {"test": preds}
    t.assert_close(mean([x], "test"), expected_mean)
//  [python]
def test_mean(preds, expected_mean):
    x = {"test": preds}
    torch.testing.assert_close(mean([x], "test"), expected_mean)
//  [python]
def test_accuracy(preds, target, expected_accuracy):
    t.assert_close(accuracy(preds, target), expected_accuracy)
//  [python]
def test_accuracy(preds, target, expected_accuracy):
    torch.testing.assert_close(accuracy(preds, target), expected_accuracy)
//  [python]
def test_auc_1d_input_no_reorder(self) -> None:
        x = torch.ones(3)
        y = torch.ones(1, 3)

        sklearn_x, sklearn_x_idx = x.unsqueeze(0).sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y)
        torch.testing.assert_allclose(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auc_1d_input_no_reorder(self) -> None:
        x = torch.ones(3)
        y = torch.ones(1, 3)

        sklearn_x, sklearn_x_idx = x.unsqueeze(0).sort(dim=1)
        sklearn_y = y.gather(1, sklearn_x_idx)
        sklearn_auc_result = torch.tensor(
            [sklearn_auc(m, n) for n, m in zip(sklearn_y, sklearn_x)],
            dtype=torch.float32,
        )
        torcheval_auc_result = auc(x, y)
        torch.testing.assert_close(
            torcheval_auc_result,
            sklearn_auc_result,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_precision_at_k(output, target, expected_precision_at_k):
    t.assert_close(precision_at_k(output, target), expected_precision_at_k)
//  [python]
def test_precision_at_k(output, target, expected_precision_at_k):
    torch.testing.assert_close(precision_at_k(output, target), expected_precision_at_k)
//  [python]
def test_feature_scalar(self):
        field = "test_feat"
        input_tensor = torch.tensor([-2.5, 0.0, 1.0], dtype=torch.float32)
        expected_tensor = torch.tensor([0.0, 2.0, 2.8], dtype=torch.float32)
        processor = FeatureScaler(field, add_value=2.5, div_value=1.25)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_feature_scalar(self):
        field = "test_feat"
        input_tensor = torch.tensor([-2.5, 0.0, 1.0], dtype=torch.float32)
        expected_tensor = torch.tensor([0.0, 2.0, 2.8], dtype=torch.float32)
        processor = FeatureScaler(field, add_value=2.5, div_value=1.25)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def test_log_compression(self):
        field = "test_feat"

        input_tensor = torch.tensor([-0.5, 0.0, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([np.log(0.5), 0.0, np.log(3.0)], dtype=torch.float32)
        processor = LogCompression(field)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_log_compression(self):
        field = "test_feat"

        input_tensor = torch.tensor([-0.5, 0.0, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([np.log(0.5), 0.0, np.log(3.0)], dtype=torch.float32)
        processor = LogCompression(field)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def compare_decoder(
    bsz: int,
    seq_len: int,
    embed_dim: int,
    num_heads: int,
    num_kv_heads: int,
    max_seq_len: int,
    vocab_size: int,
    num_layers: int,
) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(16)

    # generate input tensor used by both implementations
    x_input = torch.randint(low=0, high=vocab_size, size=(bsz, seq_len))

    # reference implementation; initialize with constant to compare outputs
    decoder_ref = Transformer(
        vocab_size=vocab_size,
        dim=embed_dim,
        n_layers=num_layers,
        n_heads=num_heads,
        max_seq_len=max_seq_len,
        n_kv_heads=num_kv_heads,
    )
    init_weights_with_constant(decoder_ref, constant=0.2)

    with torch.no_grad():
        output_ref = decoder_ref(x_input)

    # current implementation; initialize with constant to compare outputs
    decoder = llama2(
        vocab_size=vocab_size,
        embed_dim=embed_dim,
        num_layers=num_layers,
        num_heads=num_heads,
        max_seq_len=max_seq_len,
        num_kv_heads=num_kv_heads,
    )
    init_weights_with_constant(decoder, constant=0.2)

    with torch.no_grad():
        output = decoder(x_input)

    # value: tensor(20.4800)
    print(output.mean())

    torch.testing.assert_allclose(output_ref, output, atol=1e-6, rtol=1e-6)
//  [python]
def compare_decoder(
    bsz: int,
    seq_len: int,
    embed_dim: int,
    num_heads: int,
    num_kv_heads: int,
    max_seq_len: int,
    vocab_size: int,
    num_layers: int,
) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(16)

    # generate input tensor used by both implementations
    x_input = torch.randint(low=0, high=vocab_size, size=(bsz, seq_len))

    # reference implementation; initialize with constant to compare outputs
    decoder_ref = Transformer(
        vocab_size=vocab_size,
        dim=embed_dim,
        n_layers=num_layers,
        n_heads=num_heads,
        max_seq_len=max_seq_len,
        n_kv_heads=num_kv_heads,
    )
    init_weights_with_constant(decoder_ref, constant=0.2)

    with torch.no_grad():
        output_ref = decoder_ref(x_input)

    # current implementation; initialize with constant to compare outputs
    decoder = llama2(
        vocab_size=vocab_size,
        embed_dim=embed_dim,
        num_layers=num_layers,
        num_heads=num_heads,
        max_seq_len=max_seq_len,
        num_kv_heads=num_kv_heads,
    )
    init_weights_with_constant(decoder, constant=0.2)

    with torch.no_grad():
        output = decoder(x_input)

    # value: tensor(20.4800)
    print(output.mean())

    torch.testing.assert_close(output_ref, output, atol=1e-6, rtol=1e-6)
//  [python]
def test_log_compression_clamp(self):
        field = "test_feat"

        input_tensor = torch.tensor([0.1, 1.0, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([np.log(0.5), 0.0, np.log(2.0)], dtype=torch.float32)
        processor = LogCompression(field, log_zero_guard_type="clamp", log_zero_guard_value=0.5)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_log_compression_clamp(self):
        field = "test_feat"

        input_tensor = torch.tensor([0.1, 1.0, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([np.log(0.5), 0.0, np.log(2.0)], dtype=torch.float32)
        processor = LogCompression(field, log_zero_guard_type="clamp", log_zero_guard_value=0.5)

        training_example = {field: input_tensor}
        processor.process(training_example)
        output_tensor = training_example[field]

        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_normalization(self):
        field = "test_feat"
        filename = "stats.json"
        stat_dict = {"default": {"test_feat_mean": 1.5, "test_feat_std": 0.5}}

        input_tensor = torch.tensor([0.0, 1.5, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([-3.0, 0.0, 1.0], dtype=torch.float32)
        training_example = {field: input_tensor}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceNormalization(field, stats_path=stat_dict_filepath, mask_field=None)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_normalization(self):
        field = "test_feat"
        filename = "stats.json"
        stat_dict = {"default": {"test_feat_mean": 1.5, "test_feat_std": 0.5}}

        input_tensor = torch.tensor([0.0, 1.5, 2.0], dtype=torch.float32)
        expected_tensor = torch.tensor([-3.0, 0.0, 1.0], dtype=torch.float32)
        training_example = {field: input_tensor}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceNormalization(field, stats_path=stat_dict_filepath, mask_field=None)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_normalization_masked(self):
        field = "test_feat"
        mask_field = "mask"
        filename = "stats.json"
        stat_dict = {"default": {"test_feat_mean": 1.0, "test_feat_std": 0.5}}

        input_tensor = torch.tensor([2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
        input_mask = torch.tensor([True, False, False, True], dtype=torch.bool)
        expected_tensor = torch.tensor([2.0, 0.0, 0.0, 8.0], dtype=torch.float32)
        training_example = {field: input_tensor, mask_field: input_mask}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceNormalization(field, stats_path=stat_dict_filepath, mask_field=mask_field)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_normalization_masked(self):
        field = "test_feat"
        mask_field = "mask"
        filename = "stats.json"
        stat_dict = {"default": {"test_feat_mean": 1.0, "test_feat_std": 0.5}}

        input_tensor = torch.tensor([2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
        input_mask = torch.tensor([True, False, False, True], dtype=torch.bool)
        expected_tensor = torch.tensor([2.0, 0.0, 0.0, 8.0], dtype=torch.float32)
        training_example = {field: input_tensor, mask_field: input_mask}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceNormalization(field, stats_path=stat_dict_filepath, mask_field=mask_field)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_speaker_normalization_masked(self):
        field = "test_feat"
        mask_field = "test_mask"
        filename = "stats.json"
        stat_dict = {"steve": {"test_feat_mean": -1.0, "test_feat_std": 2.0}}

        input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
        input_mask = torch.tensor([False, True, False, True], dtype=torch.bool)
        expected_tensor = torch.tensor([0.0, 1.5, 0.0, 2.5], dtype=torch.float32)

        training_example = {field: input_tensor, "speaker": "steve", mask_field: input_mask}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceSpeakerNormalization(field, stats_path=stat_dict_filepath, mask_field=mask_field)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_allclose(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_speaker_normalization_masked(self):
        field = "test_feat"
        mask_field = "test_mask"
        filename = "stats.json"
        stat_dict = {"steve": {"test_feat_mean": -1.0, "test_feat_std": 2.0}}

        input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
        input_mask = torch.tensor([False, True, False, True], dtype=torch.bool)
        expected_tensor = torch.tensor([0.0, 1.5, 0.0, 2.5], dtype=torch.float32)

        training_example = {field: input_tensor, "speaker": "steve", mask_field: input_mask}

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceSpeakerNormalization(field, stats_path=stat_dict_filepath, mask_field=mask_field)
            processor.process(training_example)

        output_tensor = training_example[field]
        torch.testing.assert_close(output_tensor, expected_tensor)
//  [python]
def test_mean_variance_speaker_normalization(self):
        field = "pitch"
        filename = "stats.json"
        stat_dict = {
            "default": {"pitch_mean": 1.5, "pitch_std": 0.5},
            "speaker1": {"pitch_mean": 0.5, "pitch_std": 1.0},
            "speaker2": {"pitch_mean": 0.0, "pitch_std": 2.0},
        }

        input_tensor = torch.tensor([0.0, 1.0], dtype=torch.float32)

        training_example1 = {field: input_tensor, "speaker": "speaker1"}
        training_example2 = {field: input_tensor, "speaker": "speaker2"}
        training_example3 = {field: input_tensor, "speaker": "unknown"}
        expected_tensor1 = torch.tensor([-0.5, 0.5], dtype=torch.float32)
        expected_tensor2 = torch.tensor([0.0, 0.5], dtype=torch.float32)
        expected_tensor3 = torch.tensor([-3.0, -1.0], dtype=torch.float32)

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceSpeakerNormalization(
                field, stats_path=stat_dict_filepath, mask_field=None, fallback_to_default=True
            )
            processor.process(training_example1)
            processor.process(training_example2)
            processor.process(training_example3)

        output_tensor1 = training_example1[field]
        output_tensor2 = training_example2[field]
        output_tensor3 = training_example3[field]
        torch.testing.assert_allclose(output_tensor1, expected_tensor1)
        torch.testing.assert_allclose(output_tensor2, expected_tensor2)
        torch.testing.assert_allclose(output_tensor3, expected_tensor3)
//  [python]
def test_mean_variance_speaker_normalization(self):
        field = "pitch"
        filename = "stats.json"
        stat_dict = {
            "default": {"pitch_mean": 1.5, "pitch_std": 0.5},
            "speaker1": {"pitch_mean": 0.5, "pitch_std": 1.0},
            "speaker2": {"pitch_mean": 0.0, "pitch_std": 2.0},
        }

        input_tensor = torch.tensor([0.0, 1.0], dtype=torch.float32)

        training_example1 = {field: input_tensor, "speaker": "speaker1"}
        training_example2 = {field: input_tensor, "speaker": "speaker2"}
        training_example3 = {field: input_tensor, "speaker": "unknown"}
        expected_tensor1 = torch.tensor([-0.5, 0.5], dtype=torch.float32)
        expected_tensor2 = torch.tensor([0.0, 0.5], dtype=torch.float32)
        expected_tensor3 = torch.tensor([-3.0, -1.0], dtype=torch.float32)

        with self._write_test_dict(stat_dict, filename=filename) as stat_dict_filepath:
            processor = MeanVarianceSpeakerNormalization(
                field, stats_path=stat_dict_filepath, mask_field=None, fallback_to_default=True
            )
            processor.process(training_example1)
            processor.process(training_example2)
            processor.process(training_example3)

        output_tensor1 = training_example1[field]
        output_tensor2 = training_example2[field]
        output_tensor3 = training_example3[field]
        torch.testing.assert_close(output_tensor1, expected_tensor1)
        torch.testing.assert_close(output_tensor2, expected_tensor2)
        torch.testing.assert_close(output_tensor3, expected_tensor3)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    torch.cuda.synchronize()

    # Sets atol and rtol to looser tolerances than assert_close's defaults
    atol: Number = 1e-1
    rtol: Number = 1.3e-6
    t.assert_close(thunder_result, torch_result, equal_nan=True, atol=atol, rtol=rtol)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    torch.cuda.synchronize()

    # Sets atol and rtol to looser tolerances than assert_close's defaults
    atol: Number = 1e-1
    rtol: Number = 1.3e-6
    assert_close(thunder_result, torch_result, equal_nan=True, atol=atol, rtol=rtol)
//  [python]
def test_reorder_batched_ad_lengths_cpu(
        self,
        B: int,
        T: int,
        L: int,
        A: int,
        Dtype: torch.dtype,
        broadcast_lengths: bool,
    ) -> None:
        if broadcast_lengths:
            cat_ad_lengths = (
                torch.cat([torch.tensor([L for _ in range(T)]) for _ in range(B)], 0)
                .int()
                .to(Dtype)
            )
            cat_ad_lengths_broadcasted = cat_ad_lengths.tile([A])
        else:
            cat_ad_lengths = (
                torch.cat(
                    [torch.tensor([L for _ in range(T * A)]) for _ in range(B)], 0
                )
                .int()
                .to(Dtype)
            )
            cat_ad_lengths_broadcasted = cat_ad_lengths
        batch_offsets = torch.tensor([A * b for b in range(B + 1)]).int()
        num_ads_in_batch = B * A
        reordered_batched_ad_lengths = torch.ops.fbgemm.reorder_batched_ad_lengths(
            cat_ad_lengths, batch_offsets, num_ads_in_batch, broadcast_lengths
        )
        torch.testing.assert_allclose(
            cat_ad_lengths_broadcasted, reordered_batched_ad_lengths
        )
//  [python]
def test_reorder_batched_ad_lengths_cpu(
        self,
        B: int,
        T: int,
        L: int,
        A: int,
        Dtype: torch.dtype,
        broadcast_lengths: bool,
    ) -> None:
        if broadcast_lengths:
            cat_ad_lengths = (
                torch.cat([torch.tensor([L for _ in range(T)]) for _ in range(B)], 0)
                .int()
                .to(Dtype)
            )
            cat_ad_lengths_broadcasted = cat_ad_lengths.tile([A])
        else:
            cat_ad_lengths = (
                torch.cat(
                    [torch.tensor([L for _ in range(T * A)]) for _ in range(B)], 0
                )
                .int()
                .to(Dtype)
            )
            cat_ad_lengths_broadcasted = cat_ad_lengths
        batch_offsets = torch.tensor([A * b for b in range(B + 1)]).int()
        num_ads_in_batch = B * A
        reordered_batched_ad_lengths = torch.ops.fbgemm.reorder_batched_ad_lengths(
            cat_ad_lengths, batch_offsets, num_ads_in_batch, broadcast_lengths
        )
        torch.testing.assert_close(
            cat_ad_lengths_broadcasted, reordered_batched_ad_lengths
        )