jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_fused_obs_fq_module(self, device):
        # Set up the parameters
        x = torch.randn(5, 5, device=device)
        running_min_op = torch.tensor(float("inf"), device=device)
        running_max_op = torch.tensor(float("-inf"), device=device)
        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        # Run the forward on the Module
        mod = FusedMovingAvgObsFakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod)
        torch.ao.quantization.enable_observer(mod)
        mod.to(device)
        out = mod(x)

        # Run the operator directly
        pt_op = torch.fused_moving_avg_obs_fake_quant

        out_ref = pt_op(
            x,
            mod.observer_enabled,
            mod.fake_quant_enabled,
            running_min_op,
            running_max_op,
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )

        # Compare params with reference
        torch.testing.assert_allclose(out, out_ref)
        torch.testing.assert_allclose(
            running_min_op, mod.activation_post_process.min_val
        )
        torch.testing.assert_allclose(
            running_max_op, mod.activation_post_process.max_val
        )
//  [python]
def test_fused_obs_fq_module(self, device):
        # Set up the parameters
        x = torch.randn(5, 5, device=device)
        running_min_op = torch.tensor(float("inf"), device=device)
        running_max_op = torch.tensor(float("-inf"), device=device)
        avg_const = 0.01
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        # Run the forward on the Module
        mod = FusedMovingAvgObsFakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod)
        torch.ao.quantization.enable_observer(mod)
        mod.to(device)
        out = mod(x)

        # Run the operator directly
        pt_op = torch.fused_moving_avg_obs_fake_quant

        out_ref = pt_op(
            x,
            mod.observer_enabled,
            mod.fake_quant_enabled,
            running_min_op,
            running_max_op,
            scale,
            zero_point,
            avg_const,
            0,
            255,
            0,
            False,
        )

        # Compare params with reference
        torch.testing.assert_close(out, out_ref)
        torch.testing.assert_close(
            running_min_op, mod.activation_post_process.min_val
        )
        torch.testing.assert_close(
            running_max_op, mod.activation_post_process.max_val
        )
//  [python]
def test_masks_box(self):
        def masks_box_check(masks, expected, atol=1e-4):
            out = ops.masks_to_boxes(masks)
            assert out.dtype == torch.float
            torch.testing.assert_allclose(out, expected, rtol=0.0, check_dtype=True, atol=atol)

        # Check for int type boxes.
        def _get_image():
            assets_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
            mask_path = os.path.join(assets_directory, "masks.tiff")
            image = Image.open(mask_path)
            return image

        def _create_masks(image, masks):
            for index in range(image.n_frames):
                image.seek(index)
                frame = np.array(image)
                masks[index] = torch.tensor(frame)

            return masks

        expected = torch.tensor(
            [
                [127, 2, 165, 40],
                [2, 50, 44, 92],
                [56, 63, 98, 100],
                [139, 68, 175, 104],
                [160, 112, 198, 145],
                [49, 138, 99, 182],
                [108, 148, 152, 213],
            ],
            dtype=torch.float,
        )

        image = _get_image()
        for dtype in [torch.float16, torch.float32, torch.float64]:
            masks = torch.zeros((image.n_frames, image.height, image.width), dtype=dtype)
            masks = _create_masks(image, masks)
            masks_box_check(masks, expected)
//  [python]
def test_masks_box(self):
        def masks_box_check(masks, expected, atol=1e-4):
            out = ops.masks_to_boxes(masks)
            assert out.dtype == torch.float
            torch.testing.assert_close(out, expected, rtol=0.0, check_dtype=True, atol=atol)

        # Check for int type boxes.
        def _get_image():
            assets_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
            mask_path = os.path.join(assets_directory, "masks.tiff")
            image = Image.open(mask_path)
            return image

        def _create_masks(image, masks):
            for index in range(image.n_frames):
                image.seek(index)
                frame = np.array(image)
                masks[index] = torch.tensor(frame)

            return masks

        expected = torch.tensor(
            [
                [127, 2, 165, 40],
                [2, 50, 44, 92],
                [56, 63, 98, 100],
                [139, 68, 175, 104],
                [160, 112, 198, 145],
                [49, 138, 99, 182],
                [108, 148, 152, 213],
            ],
            dtype=torch.float,
        )

        image = _get_image()
        for dtype in [torch.float16, torch.float32, torch.float64]:
            masks = torch.zeros((image.n_frames, image.height, image.width), dtype=dtype)
            masks = _create_masks(image, masks)
            masks_box_check(masks, expected)
//  [python]
def test_no_effect_on_training(self, device, world_size, use_fsdp):
        torch.manual_seed(0)

        self._check_test_params(device, world_size, use_fsdp)

        max_duration = '6ba'

        trainer_ref = self._create_trainer(device, max_duration, use_fsdp)
        trainer_ref.fit()

        trainer_generate = self._create_trainer(
            device,
            max_duration,
            use_fsdp,
            generate_cb=Generate(prompts=['a', 'bc', 'defg'], interval=f'2ba', batch_size=2, max_new_tokens=5),
        )
        trainer_generate.fit()

        model_ref = trainer_ref.state.model
        model_generate = trainer_generate.state.model

        # Assert that the models should be equivalent
        assert model_ref is not model_generate, 'Same model should not be compared.'
        for param1, param2 in zip(model_ref.parameters(), model_generate.parameters()):
            torch.testing.assert_allclose(param1, param2)
//  [python]
def test_no_effect_on_training(self, device, world_size, use_fsdp):
        torch.manual_seed(0)

        self._check_test_params(device, world_size, use_fsdp)

        max_duration = '6ba'

        trainer_ref = self._create_trainer(device, max_duration, use_fsdp)
        trainer_ref.fit()

        trainer_generate = self._create_trainer(
            device,
            max_duration,
            use_fsdp,
            generate_cb=Generate(prompts=['a', 'bc', 'defg'], interval=f'2ba', batch_size=2, max_new_tokens=5),
        )
        trainer_generate.fit()

        model_ref = trainer_ref.state.model
        model_generate = trainer_generate.state.model

        # Assert that the models should be equivalent
        assert model_ref is not model_generate, 'Same model should not be compared.'
        for param1, param2 in zip(model_ref.parameters(), model_generate.parameters()):
            torch.testing.assert_close(param1, param2)
//  [python]
def test_fused_obs_fq_moving_avg_module(self, device):
        # Set up the parameters
        running_min_op = torch.tensor(float("inf"), device=device)
        running_max_op = torch.tensor(float("-inf"), device=device)
        avg_const = 0.001
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        mod = FusedMovingAvgObsFakeQuantize(averaging_constant=0.001)
        mod.to(device)
        mod.observer_enabled[0] = 0
        mod.fake_quant_enabled[0] = 0

        for i in range(10):
            x = torch.randn(5, 5, device=device)
            if i > 2:
                mod.observer_enabled[0] = 1
            if i > 4:
                mod.fake_quant_enabled[0] = 1
            # Run the forward on the Module
            out = mod(x)

            # Run the operator directly
            pt_op = torch.fused_moving_avg_obs_fake_quant

            out_ref = pt_op(
                x,
                mod.observer_enabled,
                mod.fake_quant_enabled,
                running_min_op,
                running_max_op,
                scale,
                zero_point,
                avg_const,
                0,
                255,
                0,
                False,
            )

            # Compare params with reference
            t.assert_close(out, out_ref)
            t.assert_close(
                running_min_op, mod.activation_post_process.min_val
            )
            t.assert_close(
                running_max_op, mod.activation_post_process.max_val
            )
//  [python]
def test_fused_obs_fq_moving_avg_module(self, device):
        # Set up the parameters
        running_min_op = torch.tensor(float("inf"), device=device)
        running_max_op = torch.tensor(float("-inf"), device=device)
        avg_const = 0.001
        scale = torch.tensor([1.0], device=device)
        zero_point = torch.tensor([0], dtype=torch.int, device=device)

        mod = FusedMovingAvgObsFakeQuantize(averaging_constant=0.001)
        mod.to(device)
        mod.observer_enabled[0] = 0
        mod.fake_quant_enabled[0] = 0

        for i in range(10):
            x = torch.randn(5, 5, device=device)
            if i > 2:
                mod.observer_enabled[0] = 1
            if i > 4:
                mod.fake_quant_enabled[0] = 1
            # Run the forward on the Module
            out = mod(x)

            # Run the operator directly
            pt_op = torch.fused_moving_avg_obs_fake_quant

            out_ref = pt_op(
                x,
                mod.observer_enabled,
                mod.fake_quant_enabled,
                running_min_op,
                running_max_op,
                scale,
                zero_point,
                avg_const,
                0,
                255,
                0,
                False,
            )

            # Compare params with reference
            torch.testing.assert_close(out, out_ref)
            torch.testing.assert_close(
                running_min_op, mod.activation_post_process.min_val
            )
            torch.testing.assert_close(
                running_max_op, mod.activation_post_process.max_val
            )
//  [python]
def check_close(a, b, atol=1e-5, rtol=1e-5):
    try:
        t.assert_close(a, b, atol=atol, rtol=rtol)
    except AssertionError:
        traceback.print_exc()
        diff = torch.abs(a - b).detach()
        print(f"{diff.max()=} {diff.mean()=}")
        import ipdb

        ipdb.set_trace()
//  [python]
def check_close(a, b, atol=1e-5, rtol=1e-5):
    try:
        torch.testing.assert_close(a, b, atol=atol, rtol=rtol)
    except AssertionError:
        traceback.print_exc()
        diff = torch.abs(a - b).detach()
        print(f"{diff.max()=} {diff.mean()=}")
        import ipdb

        ipdb.set_trace()
//  [python]
def test_op(M, N, dtype, mode, device):
    capability = torch.cuda.get_device_capability()
    if capability[0] < 8 and dtype == "bfloat16":
        pytest.skip("Only test bfloat16 on devices with sm >= 80")
    dtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}[dtype]
    # create inputs
    x = torch.randn(M, N, dtype=dtype, device=device, requires_grad=True)
    idx = 4 + torch.ones(M, dtype=torch.int64, device=device)
    # forward pass
    tt_y = triton.ops.cross_entropy(x, idx)
    th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
    if mode == 'forward':
        torch.testing.assert_allclose(th_y, tt_y)
    # backward pass
    elif mode == 'backward':
        dy = torch.randn_like(tt_y)
        # triton backward
        tt_y.backward(dy)
        tt_dx = x.grad.clone()
        # torch backward
        x.grad = None
        th_y.backward(dy)
        th_dx = x.grad.clone()
        if dtype == torch.float16:
            torch.testing.assert_allclose(th_dx, tt_dx, rtol=0.001, atol=0.001)
        else:
            torch.testing.assert_allclose(th_dx, tt_dx)
//  [python]
def test_op(M, N, dtype, mode, device):
    capability = torch.cuda.get_device_capability()
    if capability[0] < 8 and dtype == "bfloat16":
        pytest.skip("Only test bfloat16 on devices with sm >= 80")
    dtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}[dtype]
    # create inputs
    x = torch.randn(M, N, dtype=dtype, device=device, requires_grad=True)
    idx = 4 + torch.ones(M, dtype=torch.int64, device=device)
    # forward pass
    tt_y = triton.ops.cross_entropy(x, idx)
    th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
    if mode == 'forward':
        torch.testing.assert_close(th_y, tt_y)
    # backward pass
    elif mode == 'backward':
        dy = torch.randn_like(tt_y)
        # triton backward
        tt_y.backward(dy)
        tt_dx = x.grad.clone()
        # torch backward
        x.grad = None
        th_y.backward(dy)
        th_dx = x.grad.clone()
        if dtype == torch.float16:
            torch.testing.assert_close(th_dx, tt_dx, rtol=0.001, atol=0.001)
        else:
            torch.testing.assert_close(th_dx, tt_dx)
//  [python]
def test_forward_lstm_cell(self, sizes, _seed):
        p = sizes
        ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False)
        test_data = LSTMTestData(*clone_test_data(ref_data))

        ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(ref_rnn, ref_data)
        test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(test_rnn, test_data)

        for i in range(p.seq_length):
            ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
            test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
            for ref, test in list(zip(ref_result, test_result)):
                torch.allclose(test, ref)
//  [python]
def test_forward_lstm_cell(self, sizes, _seed):
        p = sizes
        ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False)
        test_data = LSTMTestData(*clone_test_data(ref_data))

        ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(ref_rnn, ref_data)
        test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(test_rnn, test_data)

        for i in range(p.seq_length):
            ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
            test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
            for ref, test in list(zip(ref_result, test_result)):
                torch.testing.assert_close(test, ref)
//  [python]
def _test_full(
        self,
        shape,
        fill_value,
        dtype="float16",
        test_name="full",
    ) -> None:
        Y = ops.full()(shape, fill_value, dtype)
        Y._attrs["name"] = "Y"

        if not isinstance(shape, list):
            shape = [shape]

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        Z = ops.elementwise(FuncEnum.ADD)(X, Y)
        Z._attrs["name"] = "Z"
        Z._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Z, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        if isinstance(shape[0], IntVar):
            shapes = [[val] + shape[1:] for val in shape[0]._attrs["values"]]
        else:
            shapes = [shape]

        for shape in shapes:
            x_pt = get_random_torch_tensor(shape, dtype=dtype)
            z_pt = x_pt + fill_value

            z = torch.empty_like(z_pt)

            module.run_with_tensors([x_pt], [z])
            t.assert_allclose(z, z_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_full(
        self,
        shape,
        fill_value,
        dtype="float16",
        test_name="full",
    ) -> None:
        Y = ops.full()(shape, fill_value, dtype)
        Y._attrs["name"] = "Y"

        if not isinstance(shape, list):
            shape = [shape]

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        Z = ops.elementwise(FuncEnum.ADD)(X, Y)
        Z._attrs["name"] = "Z"
        Z._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Z, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        if isinstance(shape[0], IntVar):
            shapes = [[val] + shape[1:] for val in shape[0]._attrs["values"]]
        else:
            shapes = [shape]

        for shape in shapes:
            x_pt = get_random_torch_tensor(shape, dtype=dtype)
            z_pt = x_pt + fill_value

            z = torch.empty_like(z_pt)

            module.run_with_tensors([x_pt], [z])
            torch.testing.assert_close(z, z_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_const_full_with_copy(self, dtype="float16"):
        shape = [IntVar([1, 128]), 10]
        full = ops.full()(shape, 1.0, dtype)
        add = ops.elementwise(FuncEnum.ADD)(full, 1.0)
        Y = ops.flatten()(add)

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        flattened_X = ops.flatten()(X)
        Z = ops.elementwise(FuncEnum.ADD)(flattened_X, Y)
        Z._attrs["name"] = "Z"
        Z._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Z, target, "./tmp", "const_full_with_copy")

        if isinstance(shape[0], IntVar):
            shapes = [[val] + shape[1:] for val in shape[0]._attrs["values"]]
        else:
            shapes = [shape]

        for shape in shapes:
            x_pt = get_random_torch_tensor(shape, dtype=dtype)
            tmp_pt = x_pt + 1.0 + 1.0
            z_pt = torch.flatten(tmp_pt)

            z = torch.empty_like(z_pt)

            module.run_with_tensors([x_pt], [z])
            t.assert_close(z, z_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_const_full_with_copy(self, dtype="float16"):
        shape = [IntVar([1, 128]), 10]
        full = ops.full()(shape, 1.0, dtype)
        add = ops.elementwise(FuncEnum.ADD)(full, 1.0)
        Y = ops.flatten()(add)

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        flattened_X = ops.flatten()(X)
        Z = ops.elementwise(FuncEnum.ADD)(flattened_X, Y)
        Z._attrs["name"] = "Z"
        Z._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Z, target, "./tmp", "const_full_with_copy")

        if isinstance(shape[0], IntVar):
            shapes = [[val] + shape[1:] for val in shape[0]._attrs["values"]]
        else:
            shapes = [shape]

        for shape in shapes:
            x_pt = get_random_torch_tensor(shape, dtype=dtype)
            tmp_pt = x_pt + 1.0 + 1.0
            z_pt = torch.flatten(tmp_pt)

            z = torch.empty_like(z_pt)

            module.run_with_tensors([x_pt], [z])
            torch.testing.assert_close(z, z_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_fused_softmax(self):
        grads = torch.randn_like(self.tst_inputs)
        y_ref = self.ref_inputs.view(self.sequences, self.heads, self.seq_length, self.seq_length)
        y_ref = y_ref + self.mask.unsqueeze(1).unsqueeze(2)
        y_ref = y_ref.view(self.sequences * self.heads, self.seq_length, self.seq_length)
        y_ref = F.softmax(y_ref, dim=-1)
        y_ref = torch._fused_dropout(y_ref, 1.0)

        y_tst = fast_mask_softmax_dropout_func(True, self.heads, self.tst_inputs, self.mask, True, 0.0)
        y_ref[0].backward(grads)
        y_tst.backward(grads)

        torch.testing.assert_allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(y_ref[0], y_tst, atol=1e-3, rtol=1e-3)
        torch.testing.assert_allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_fused_softmax(self):
        grads = torch.randn_like(self.tst_inputs)
        y_ref = self.ref_inputs.view(self.sequences, self.heads, self.seq_length, self.seq_length)
        y_ref = y_ref + self.mask.unsqueeze(1).unsqueeze(2)
        y_ref = y_ref.view(self.sequences * self.heads, self.seq_length, self.seq_length)
        y_ref = F.softmax(y_ref, dim=-1)
        y_ref = torch._fused_dropout(y_ref, 1.0)

        y_tst = fast_mask_softmax_dropout_func(True, self.heads, self.tst_inputs, self.mask, True, 0.0)
        y_ref[0].backward(grads)
        y_tst.backward(grads)

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(y_ref[0], y_tst, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_generate_sem_numpy(five_node_graph: torch.Tensor, five_variable_one_dim_metadata: VariablesMetadata):
    tmpfile = tempfile.mkstemp(suffix=".npy")[1]

    np.save(tmpfile, five_node_graph.numpy())
    assert os.path.exists(tmpfile)

    variables = get_variable_definitions(num_nodes=5)

    sem = generate_sem_sampler(
        variables=variables,
        graph_file=tmpfile,
        num_edges=None,
        graph_type="numpy",
        function_type="linear",
    ).sample()
    assert len(variables.variables) == 5
    assert variables == five_variable_one_dim_metadata

    samples = sem.sample(torch.Size([5]))
    assert samples["x_0"].shape == torch.Size([5, 1])
    assert samples["x_1"].shape == torch.Size([5, 1])
    assert samples["x_2"].shape == torch.Size([5, 1])
    assert samples["x_3"].shape == torch.Size([5, 1])
    assert samples["x_4"].shape == torch.Size([5, 1])

    t.assert_close(sem.graph, five_node_graph)
//  [python]
def test_generate_sem_numpy(five_node_graph: torch.Tensor, five_variable_one_dim_metadata: VariablesMetadata):
    tmpfile = tempfile.mkstemp(suffix=".npy")[1]

    np.save(tmpfile, five_node_graph.numpy())
    assert os.path.exists(tmpfile)

    variables = get_variable_definitions(num_nodes=5)

    sem = generate_sem_sampler(
        variables=variables,
        graph_file=tmpfile,
        num_edges=None,
        graph_type="numpy",
        function_type="linear",
    ).sample()
    assert len(variables.variables) == 5
    assert variables == five_variable_one_dim_metadata

    samples = sem.sample(torch.Size([5]))
    assert samples["x_0"].shape == torch.Size([5, 1])
    assert samples["x_1"].shape == torch.Size([5, 1])
    assert samples["x_2"].shape == torch.Size([5, 1])
    assert samples["x_3"].shape == torch.Size([5, 1])
    assert samples["x_4"].shape == torch.Size([5, 1])

    torch.testing.assert_close(sem.graph, five_node_graph)
//  [python]
def test_sample_counterfactuals(mixture_sem_sampler_and_shapes: tuple[SEMSampler, dict[str, torch.Size]]):
    mixture_sem_sampler, shapes_dict = mixture_sem_sampler_and_shapes
    sem = mixture_sem_sampler.sample()

    num_samples = 5
    noise = sem.sample_noise(torch.Size([num_samples]))
    observations = sem.noise_to_sample(noise)

    counterfactual = sample_counterfactual(sem, observations, noise, "x_1")

    for v, shape in shapes_dict.items():
        assert counterfactual.counterfactual_data[v].shape == torch.Size([num_samples]) + shape
    # test ordering of intervention data
    assert list(counterfactual.counterfactual_data.keys()) == ["x_0", "x_1", "x_2", "x_3", "x_4"]

    assert all(val.ndim == 1 for val in counterfactual.intervention_values.values())
    assert counterfactual.intervention_values["x_1"].shape == shapes_dict["x_1"]
    assert all(counterfactual.intervention_values["x_1"] >= observations["x_1"].min())
    assert all(counterfactual.intervention_values["x_1"] <= observations["x_1"].max())

    torch.testing.assert_allclose(counterfactual.factual_data, observations)
    torch.testing.assert_allclose(
        sem.do(counterfactual.intervention_values).noise_to_sample(noise), counterfactual.counterfactual_data
    )
//  [python]
def test_sample_counterfactuals(mixture_sem_sampler_and_shapes: tuple[SEMSampler, dict[str, torch.Size]]):
    mixture_sem_sampler, shapes_dict = mixture_sem_sampler_and_shapes
    sem = mixture_sem_sampler.sample()

    num_samples = 5
    noise = sem.sample_noise(torch.Size([num_samples]))
    observations = sem.noise_to_sample(noise)

    counterfactual = sample_counterfactual(sem, observations, noise, "x_1")

    for v, shape in shapes_dict.items():
        assert counterfactual.counterfactual_data[v].shape == torch.Size([num_samples]) + shape
    # test ordering of intervention data
    assert list(counterfactual.counterfactual_data.keys()) == ["x_0", "x_1", "x_2", "x_3", "x_4"]

    assert all(val.ndim == 1 for val in counterfactual.intervention_values.values())
    assert counterfactual.intervention_values["x_1"].shape == shapes_dict["x_1"]
    assert all(counterfactual.intervention_values["x_1"] >= observations["x_1"].min())
    assert all(counterfactual.intervention_values["x_1"] <= observations["x_1"].max())

    torch.testing.assert_close(counterfactual.factual_data, observations)
    torch.testing.assert_close(
        sem.do(counterfactual.intervention_values).noise_to_sample(noise), counterfactual.counterfactual_data
    )
//  [python]
def _check_vectorized_wrappers(vectorization, objective):
    for w in vectorization._cost_fn_wrappers:
        for cost_fn in objective.cost_functions.values():
            if cost_fn is w.cost_fn:
                w_jac, w_err = cost_fn.weighted_jacobians_error()
                t.assert_close(w._cached_error, w_err)
                for jac, exp_jac in zip(w._cached_jacobians, w_jac):
                    t.assert_close(jac, exp_jac, atol=1e-6, rtol=1e-6)
//  [python]
def _check_vectorized_wrappers(vectorization, objective):
    for w in vectorization._cost_fn_wrappers:
        for cost_fn in objective.cost_functions.values():
            if cost_fn is w.cost_fn:
                w_jac, w_err = cost_fn.weighted_jacobians_error()
                torch.testing.assert_close(w._cached_error, w_err)
                for jac, exp_jac in zip(w._cached_jacobians, w_jac):
                    torch.testing.assert_close(jac, exp_jac, atol=1e-6, rtol=1e-6)
//  [python]
def test_backward_lstm_cell(self, sizes, _seed):
        p = sizes
        ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False, is_backward=True)
        with torch.no_grad():
            test_data = LSTMTestData(*clone_test_data(ref_data))

        ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(ref_rnn, ref_data)
        test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(test_rnn, test_data)

        for i in range(p.seq_length):
            ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
            test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
            ref_result[0].sum().backward()
            test_result[0].sum().backward()
            ref_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
            ref_grads += get_grads([ref_rnn.weight_ih, ref_rnn.weight_hh, ref_rnn.bias_ih, ref_rnn.bias_hh])
            test_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
            test_grads += get_grads([test_rnn.weight_ih, test_rnn.weight_hh, test_rnn.bias_ih, test_rnn.bias_hh])
            for ref, test in list(zip(test_grads, ref_grads)):
                torch.testing.assert_allclose(test, ref)
//  [python]
def test_backward_lstm_cell(self, sizes, _seed):
        p = sizes
        ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False, is_backward=True)
        with torch.no_grad():
            test_data = LSTMTestData(*clone_test_data(ref_data))

        ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(ref_rnn, ref_data)
        test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
        TestLSTMCell.set_weights(test_rnn, test_data)

        for i in range(p.seq_length):
            ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
            test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
            ref_result[0].sum().backward()
            test_result[0].sum().backward()
            ref_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
            ref_grads += get_grads([ref_rnn.weight_ih, ref_rnn.weight_hh, ref_rnn.bias_ih, ref_rnn.bias_hh])
            test_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
            test_grads += get_grads([test_rnn.weight_ih, test_rnn.weight_hh, test_rnn.bias_ih, test_rnn.bias_hh])
            for ref, test in list(zip(test_grads, ref_grads)):
                torch.testing.assert_close(test, ref)
//  [python]
def test_masked_jacobians(batch_size, dof, num_costs, weight_cls):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    sol1 = _solve_fn_for_masked_jacobians(
        batch_size, dof, num_costs, weight_cls, True, device
    )
    sol2 = _solve_fn_for_masked_jacobians(
        batch_size, dof, num_costs, weight_cls, False, device
    )

    for i in range(num_costs):
        torch.testing.assert_allclose(sol1[f"x{i}"], sol2[f"x{i}"])
//  [python]
def test_masked_jacobians(batch_size, dof, num_costs, weight_cls):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    sol1 = _solve_fn_for_masked_jacobians(
        batch_size, dof, num_costs, weight_cls, True, device
    )
    sol2 = _solve_fn_for_masked_jacobians(
        batch_size, dof, num_costs, weight_cls, False, device
    )

    for i in range(num_costs):
        torch.testing.assert_close(sol1[f"x{i}"], sol2[f"x{i}"])
//  [python]
def test_transform_env(self, skip):
        """Tests that the built-in frame_skip and the transform lead to the same results."""
        torch.manual_seed(0)
        if skip < 0:
            with pytest.raises(
                ValueError,
                match="frame_skip should have a value greater or equal to one",
            ):
                FrameSkipTransform(skip)
            return
        else:
            fs = FrameSkipTransform(skip)
        base_env = GymEnv(PENDULUM_VERSIONED(), frame_skip=skip)
        tensordicts = TensorDict({"action": base_env.action_spec.rand((10,))}, [10])
        env = TransformedEnv(GymEnv(PENDULUM_VERSIONED()), fs)
        base_env.set_seed(0)
        env.base_env.set_seed(0)
        td1 = base_env.reset()
        td2 = env.reset()
        for key in td1.keys():
            t.assert_close(td1[key], td2[key])
        for i in range(10):
            td1 = base_env.step(tensordicts[i].clone()).flatten_keys()
            td2 = env.step(tensordicts[i].clone()).flatten_keys()
            for key in td1.keys():
                t.assert_close(td1[key], td2[key])
//  [python]
def test_transform_env(self, skip):
        """Tests that the built-in frame_skip and the transform lead to the same results."""
        torch.manual_seed(0)
        if skip < 0:
            with pytest.raises(
                ValueError,
                match="frame_skip should have a value greater or equal to one",
            ):
                FrameSkipTransform(skip)
            return
        else:
            fs = FrameSkipTransform(skip)
        base_env = GymEnv(PENDULUM_VERSIONED(), frame_skip=skip)
        tensordicts = TensorDict({"action": base_env.action_spec.rand((10,))}, [10])
        env = TransformedEnv(GymEnv(PENDULUM_VERSIONED()), fs)
        base_env.set_seed(0)
        env.base_env.set_seed(0)
        td1 = base_env.reset()
        td2 = env.reset()
        for key in td1.keys():
            torch.testing.assert_close(td1[key], td2[key])
        for i in range(10):
            td1 = base_env.step(tensordicts[i].clone()).flatten_keys()
            td2 = env.step(tensordicts[i].clone()).flatten_keys()
            for key in td1.keys():
                torch.testing.assert_close(td1[key], td2[key])
//  [python]
def test_tranform_offline_against_online(self, padding, envtype):
        torch.manual_seed(0)
        key = "observation" if envtype == "gym" else "pixels"
        env = SerialEnv(
            3,
            lambda: TransformedEnv(
                GymEnv("CartPole-v1")
                if envtype == "gym"
                else DiscreteActionConvMockEnv(),
                CatFrames(
                    dim=-3 if envtype == "conv" else -1,
                    N=5,
                    in_keys=[key],
                    out_keys=[f"{key}_cat"],
                    padding=padding,
                ),
            ),
        )
        env.set_seed(0)

        r = env.rollout(100, break_when_any_done=False)

        c = CatFrames(
            dim=-3 if envtype == "conv" else -1,
            N=5,
            in_keys=[key, ("next", key)],
            out_keys=[f"{key}_cat2", ("next", f"{key}_cat2")],
            padding=padding,
        )

        r2 = c(r)

        torch.testing.assert_allclose(r2[f"{key}_cat2"], r2[f"{key}_cat"])
        torch.testing.assert_allclose(r2["next", f"{key}_cat2"], r2["next", f"{key}_cat"])
//  [python]
def test_tranform_offline_against_online(self, padding, envtype):
        torch.manual_seed(0)
        key = "observation" if envtype == "gym" else "pixels"
        env = SerialEnv(
            3,
            lambda: TransformedEnv(
                GymEnv("CartPole-v1")
                if envtype == "gym"
                else DiscreteActionConvMockEnv(),
                CatFrames(
                    dim=-3 if envtype == "conv" else -1,
                    N=5,
                    in_keys=[key],
                    out_keys=[f"{key}_cat"],
                    padding=padding,
                ),
            ),
        )
        env.set_seed(0)

        r = env.rollout(100, break_when_any_done=False)

        c = CatFrames(
            dim=-3 if envtype == "conv" else -1,
            N=5,
            in_keys=[key, ("next", key)],
            out_keys=[f"{key}_cat2", ("next", f"{key}_cat2")],
            padding=padding,
        )

        r2 = c(r)

        torch.testing.assert_close(r2[f"{key}_cat2"], r2[f"{key}_cat"])
        torch.testing.assert_close(r2["next", f"{key}_cat2"], r2["next", f"{key}_cat"])
//  [python]
def test_op_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):
    torch.manual_seed(20)
    q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    if TORCH_HAS_FP8:
        q = q.to(torch_dtype)
        k = k.to(torch_dtype)
    sm_scale = 0.5
    dout = torch.randn_like(q, dtype=torch.float16)
    # reference implementation
    M = torch.tril(torch.ones((N_CTX, N_CTX), device="cuda"))
    p = torch.matmul(q.half(), k.transpose(2, 3).half()) * sm_scale
    if causal:
        p[:, :, M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).half()
    ref_out = torch.matmul(p, v)
    # triton implementation
    tri_out = attention(q, k, v, causal, sm_scale)
    # compare
    torch.allclose(ref_out, tri_out, atol=1e-2, rtol=1e-2)
//  [python]
def test_op_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):
    torch.manual_seed(20)
    q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0., std=0.5).requires_grad_()
    if TORCH_HAS_FP8:
        q = q.to(torch_dtype)
        k = k.to(torch_dtype)
    sm_scale = 0.5
    dout = torch.randn_like(q, dtype=torch.float16)
    # reference implementation
    M = torch.tril(torch.ones((N_CTX, N_CTX), device="cuda"))
    p = torch.matmul(q.half(), k.transpose(2, 3).half()) * sm_scale
    if causal:
        p[:, :, M == 0] = float("-inf")
    p = torch.softmax(p.float(), dim=-1).half()
    ref_out = torch.matmul(p, v)
    # triton implementation
    tri_out = attention(q, k, v, causal, sm_scale)
    # compare
    torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2)
//  [python]
def test_transform_env_clone(self, standard_normal):
        out_keys = ["stuff"]
        if standard_normal:
            scale = 1_000_000
        else:
            scale = 0.0
        env = TransformedEnv(
            ContinuousActionVecMockEnv(),
            ObservationNorm(
                loc=0.0,
                scale=scale,
                in_keys=["observation"],
                out_keys=out_keys,
                standard_normal=standard_normal,
            ),
        )
        cloned = env.transform.clone()
        env.transform.loc += 1
        env.transform.scale += 1
        t.assert_close(
            env.transform.loc, torch.ones_like(env.transform.loc)
        )
        t.assert_close(
            env.transform.scale, scale + torch.ones_like(env.transform.scale)
        )
        assert env.transform.loc == cloned.loc
        assert env.transform.scale == cloned.scale
//  [python]
def test_transform_env_clone(self, standard_normal):
        out_keys = ["stuff"]
        if standard_normal:
            scale = 1_000_000
        else:
            scale = 0.0
        env = TransformedEnv(
            ContinuousActionVecMockEnv(),
            ObservationNorm(
                loc=0.0,
                scale=scale,
                in_keys=["observation"],
                out_keys=out_keys,
                standard_normal=standard_normal,
            ),
        )
        cloned = env.transform.clone()
        env.transform.loc += 1
        env.transform.scale += 1
        torch.testing.assert_close(
            env.transform.loc, torch.ones_like(env.transform.loc)
        )
        torch.testing.assert_close(
            env.transform.scale, scale + torch.ones_like(env.transform.scale)
        )
        assert env.transform.loc == cloned.loc
        assert env.transform.scale == cloned.scale
//  [python]
def test_frame_skip_transform_unroll(self, skip):
        torch.manual_seed(0)
        if skip < 0:
            with pytest.raises(
                ValueError,
                match="frame_skip should have a value greater or equal to one",
            ):
                FrameSkipTransform(skip)
            return
        else:
            fs = FrameSkipTransform(skip)
        base_env = GymEnv(PENDULUM_VERSIONED())
        tensordicts = TensorDict({"action": base_env.action_spec.rand((10,))}, [10])
        env = TransformedEnv(GymEnv(PENDULUM_VERSIONED()), fs)
        base_env.set_seed(0)
        env.base_env.set_seed(0)
        td1 = base_env.reset()
        td2 = env.reset()
        for key in td1.keys():
            t.assert_close(td1[key], td2[key])
        for i in range(10):
            r = 0.0
            for _ in range(skip):
                td1 = base_env.step(tensordicts[i].clone()).flatten_keys(".")
                r = td1.get("next.reward") + r
            td1.set("next.reward", r)
            td2 = env.step(tensordicts[i].clone()).flatten_keys(".")
            for key in td1.keys():
                t.assert_close(td1[key], td2[key])
//  [python]
def test_frame_skip_transform_unroll(self, skip):
        torch.manual_seed(0)
        if skip < 0:
            with pytest.raises(
                ValueError,
                match="frame_skip should have a value greater or equal to one",
            ):
                FrameSkipTransform(skip)
            return
        else:
            fs = FrameSkipTransform(skip)
        base_env = GymEnv(PENDULUM_VERSIONED())
        tensordicts = TensorDict({"action": base_env.action_spec.rand((10,))}, [10])
        env = TransformedEnv(GymEnv(PENDULUM_VERSIONED()), fs)
        base_env.set_seed(0)
        env.base_env.set_seed(0)
        td1 = base_env.reset()
        td2 = env.reset()
        for key in td1.keys():
            torch.testing.assert_close(td1[key], td2[key])
        for i in range(10):
            r = 0.0
            for _ in range(skip):
                td1 = base_env.step(tensordicts[i].clone()).flatten_keys(".")
                r = td1.get("next.reward") + r
            td1.set("next.reward", r)
            td2 = env.step(tensordicts[i].clone()).flatten_keys(".")
            for key in td1.keys():
                torch.testing.assert_close(td1[key], td2[key])
//  [python]
def test_jagged_jagged_bmm(
        self,
        B: int,
        M: int,
        N: int,
        max_L: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        assume(B != 0)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        total_length = int(lengths.sum().item())
        x_values = torch.rand(
            (total_length, M), requires_grad=True, dtype=dtype, device=device
        )
        y_values = torch.rand(
            (total_length, N), requires_grad=True, dtype=dtype, device=device
        )
        output = torch.ops.fbgemm.jagged_jagged_bmm(
            x_values,
            y_values,
            offsets,
            max_L,
        )

        x_values_ref = x_values.detach().clone().requires_grad_(True)
        y_values_ref = y_values.detach().clone().requires_grad_(True)
        x_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            x_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        y_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            y_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        output_ref = torch.bmm(x_dense_ref.transpose(2, 1), y_dense_ref)

        # verify forward
        t.assert_close(output, output_ref)

        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        t.assert_close(x_values.grad, x_values_ref.grad)
        t.assert_close(y_values.grad, y_values_ref.grad)
//  [python]
def test_jagged_jagged_bmm(
        self,
        B: int,
        M: int,
        N: int,
        max_L: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        assume(B != 0)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        total_length = int(lengths.sum().item())
        x_values = torch.rand(
            (total_length, M), requires_grad=True, dtype=dtype, device=device
        )
        y_values = torch.rand(
            (total_length, N), requires_grad=True, dtype=dtype, device=device
        )
        output = torch.ops.fbgemm.jagged_jagged_bmm(
            x_values,
            y_values,
            offsets,
            max_L,
        )

        x_values_ref = x_values.detach().clone().requires_grad_(True)
        y_values_ref = y_values.detach().clone().requires_grad_(True)
        x_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            x_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        y_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            y_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        output_ref = torch.bmm(x_dense_ref.transpose(2, 1), y_dense_ref)

        # verify forward
        torch.testing.assert_close(output, output_ref)

        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        torch.testing.assert_close(x_values.grad, x_values_ref.grad)
        torch.testing.assert_close(y_values.grad, y_values_ref.grad)
//  [python]
def test_vectorized_error():
    rng = np.random.default_rng(0)
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(20):
        dim = rng.choice([1, 2])
        objective = th.Objective()
        batch_size = rng.choice(range(1, 11))

        n_vars = rng.choice([1, 10])
        vectors = [
            th.Vector(
                tensor=torch.randn(batch_size, dim, generator=generator), name=f"v{i}"
            )
            for i in range(n_vars)
        ]
        target = th.Vector(dim, name="target")
        w = th.ScaleCostWeight(torch.randn(1, generator=generator))
        for v in vectors:
            objective.add(th.Difference(v, target, w))

        se3s = [
            th.SE3(
                tensor=th.SE3.rand(batch_size, generator=generator).tensor,
                strict_checks=False,
            )
            for i in range(rng.choice([1, 10]))
        ]
        s_target = th.SE3.rand(1, generator=generator)
        ws = th.DiagonalCostWeight(torch.randn(1, 6, generator=generator))
        for s in se3s:
            objective.add(th.Difference(s, s_target, ws))

        vectorization = th.Vectorize(objective)
        objective.update_vectorization_if_needed()

        assert objective._vectorized_jacobians_iter is vectorization._cost_fn_wrappers
        _check_vectorized_wrappers(vectorization, objective)

        squared_error = torch.cat(
            [cf.weighted_error() for cf in objective.cost_functions.values()], dim=1
        )
        torch.testing.assert_allclose(squared_error, objective.error())

        # Check that calling error(also_update=False) changes the result but doesn't
        # change the vectorized wrappers
        another_squared_error = objective.error(
            {f"v{i}": torch.ones(batch_size, dim) for i in range(n_vars)},
            also_update=False,
        )
        assert not another_squared_error.allclose(squared_error)
        _check_vectorized_wrappers(vectorization, objective)

        # Just to be sure, check that objective.error() is calling
        # the vectorized error iterator.
        called = [False]

        def mock_err_iter(self):
            called[0] = True
            return self._cost_fn_wrappers  # just to have the same return type

        with mock.patch.object(
            th.Vectorize, "_get_vectorized_error_iter", mock_err_iter
        ):
            th.Vectorize(objective)
            objective.error()
            assert called[0]