jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_non_default_quant_descriptors_are_piped(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer
        from pytorch_quantization.calib import MaxCalibrator

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer(default_quant_modules_calibrator_inputs="max", default_quant_modules_calibrator_weights="max")
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def compare_model_weights(self, base_model, model):
        for name, base_param in base_model.named_parameters():
            param = model.get_parameter(name)
            if base_param.grad is None:
                continue
            if isinstance(param, DTensor):
                param = param.redistribute(param.device_mesh, [Replicate()], async_op=False)._local_tensor
            t.testing.assert_close(param, base_param, atol=2e-4, rtol=2e-4)
//  [python]
def compare_model_weights(self, base_model, model):
        for name, base_param in base_model.named_parameters():
            param = model.get_parameter(name)
            if base_param.grad is None:
                continue
            if isinstance(param, DTensor):
                param = param.redistribute(param.device_mesh, [Replicate()], async_op=False)._local_tensor
            torch.testing.assert_close(param, base_param, atol=2e-4, rtol=2e-4)
//  [python]
def test_unsloth_cross_entropy(reduction):
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import unsloth_ex

    logits = torch.randn(64, 128, device="cuda", requires_grad=True)
    labels = torch.randint(128, (64,), device="cuda")

    def foo(logits, labels):
        # this is the variant supported by unsloth.
        # if different arguments are used, the implementation would no be lowered to unsloth and instead would get
        # decomposed
        return torch.nn.functional.cross_entropy(logits, labels, reduction=reduction, ignore_index=-100)

    cfoo = thunder.jit(foo, executors=[unsloth_ex])
    actual = cfoo(logits, labels)
    trace_str = str(thunder.last_traces(cfoo)[-1])
    assert "unsloth_cross_entropy" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cfoo)[-1])
    assert "unsloth_cross_entropy_backward" in trace_str

    expected = foo(logits, labels)
    torch.testing.assert_allclose(actual, expected)

    cfoo_grad = grad(cfoo)
    actual = cfoo_grad(logits, labels)[0]
    trace_str = str(thunder.last_traces(cfoo_grad)[-1])
    assert "unsloth_cross_entropy_backward" in trace_str
    out = foo(logits, labels)
    assert logits.grad is None
    out.sum().backward()
    expected = logits.grad
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_unsloth_cross_entropy(reduction):
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import unsloth_ex

    logits = torch.randn(64, 128, device="cuda", requires_grad=True)
    labels = torch.randint(128, (64,), device="cuda")

    def foo(logits, labels):
        # this is the variant supported by unsloth.
        # if different arguments are used, the implementation would no be lowered to unsloth and instead would get
        # decomposed
        return torch.nn.functional.cross_entropy(logits, labels, reduction=reduction, ignore_index=-100)

    cfoo = thunder.jit(foo, executors=[unsloth_ex])
    actual = cfoo(logits, labels)
    trace_str = str(thunder.last_traces(cfoo)[-1])
    assert "unsloth_cross_entropy" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cfoo)[-1])
    assert "unsloth_cross_entropy_backward" in trace_str

    expected = foo(logits, labels)
    torch.testing.assert_close(actual, expected)

    cfoo_grad = grad(cfoo)
    actual = cfoo_grad(logits, labels)[0]
    trace_str = str(thunder.last_traces(cfoo_grad)[-1])
    assert "unsloth_cross_entropy_backward" in trace_str
    out = foo(logits, labels)
    assert logits.grad is None
    out.sum().backward()
    expected = logits.grad
    torch.testing.assert_close(actual, expected)
//  [python]
def test_unsloth_swiglu():
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import ThunderLLaMAMLP, unsloth_ex
    from litgpt import Config
    from litgpt.model import LLaMAMLP

    config = Config.from_name("Llama-2-7b-hf")
    with torch.device("cuda"):
        x = torch.randn(2, 16, config.n_embd, requires_grad=True)
        mlp = LLaMAMLP(config)
    # monkeypatching was successful
    assert isinstance(mlp, ThunderLLaMAMLP)

    cmlp = thunder.jit(mlp, executors=[unsloth_ex])
    actual = cmlp(x)
    trace_str = str(thunder.last_traces(cmlp)[-1])
    assert "unsloth_swiglu" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cmlp)[-1])
    assert "unsloth_swiglu_backward" in trace_str

    expected = mlp(x)
    torch.testing.assert_allclose(actual, expected)

    cmlp_grad = grad(cmlp)
    actual = cmlp_grad(x)[0]
    trace_str = str(thunder.last_traces(cmlp_grad)[-1])
    assert "unsloth_swiglu_backward" in trace_str
    out = mlp(x)
    assert x.grad is None
    out.sum().backward()
    expected = x.grad
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_unsloth_swiglu():
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import ThunderLLaMAMLP, unsloth_ex
    from litgpt import Config
    from litgpt.model import LLaMAMLP

    config = Config.from_name("Llama-2-7b-hf")
    with torch.device("cuda"):
        x = torch.randn(2, 16, config.n_embd, requires_grad=True)
        mlp = LLaMAMLP(config)
    # monkeypatching was successful
    assert isinstance(mlp, ThunderLLaMAMLP)

    cmlp = thunder.jit(mlp, executors=[unsloth_ex])
    actual = cmlp(x)
    trace_str = str(thunder.last_traces(cmlp)[-1])
    assert "unsloth_swiglu" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cmlp)[-1])
    assert "unsloth_swiglu_backward" in trace_str

    expected = mlp(x)
    torch.testing.assert_close(actual, expected)

    cmlp_grad = grad(cmlp)
    actual = cmlp_grad(x)[0]
    trace_str = str(thunder.last_traces(cmlp_grad)[-1])
    assert "unsloth_swiglu_backward" in trace_str
    out = mlp(x)
    assert x.grad is None
    out.sum().backward()
    expected = x.grad
    torch.testing.assert_close(actual, expected)
//  [python]
def test_unsloth_rope():
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import unsloth_ex

    B, nh, T, hs = 2, 32, 64, 16
    cos, sin = build_rope_cache(T, hs, device="cuda")
    q = torch.rand((B, nh, T, hs), device="cuda", requires_grad=True)

    def foo(x, cos, sin):
        return apply_rope(x, cos, sin)

    cfoo = thunder.jit(foo, executors=[unsloth_ex])
    actual = cfoo(q, cos, sin)
    trace_str = str(thunder.last_traces(cfoo)[-1])
    assert "unsloth_apply_rope" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cfoo)[-1])
    assert "unsloth_apply_rope_backward" in trace_str

    expected = foo(q, cos, sin)
    torch.testing.assert_allclose(actual, expected)

    cfoo_grad = grad(cfoo)
    actual = cfoo_grad(q, cos, sin)[0]
    trace_str = str(thunder.last_traces(cfoo_grad)[-1])
    assert "unsloth_apply_rope_backward" in trace_str
    out = foo(q, cos, sin)
    assert q.grad is None
    out.sum().backward()
    expected = q.grad
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_unsloth_rope():
    import thunder
    from thunder.core.transforms import grad

    from extensions.thunder.unsloth.executor import unsloth_ex

    B, nh, T, hs = 2, 32, 64, 16
    cos, sin = build_rope_cache(T, hs, device="cuda")
    q = torch.rand((B, nh, T, hs), device="cuda", requires_grad=True)

    def foo(x, cos, sin):
        return apply_rope(x, cos, sin)

    cfoo = thunder.jit(foo, executors=[unsloth_ex])
    actual = cfoo(q, cos, sin)
    trace_str = str(thunder.last_traces(cfoo)[-1])
    assert "unsloth_apply_rope" in trace_str and "backward" not in trace_str
    trace_str = str(thunder.last_backward_traces(cfoo)[-1])
    assert "unsloth_apply_rope_backward" in trace_str

    expected = foo(q, cos, sin)
    torch.testing.assert_close(actual, expected)

    cfoo_grad = grad(cfoo)
    actual = cfoo_grad(q, cos, sin)[0]
    trace_str = str(thunder.last_traces(cfoo_grad)[-1])
    assert "unsloth_apply_rope_backward" in trace_str
    out = foo(q, cos, sin)
    assert q.grad is None
    out.sum().backward()
    expected = q.grad
    torch.testing.assert_close(actual, expected)
//  [python]
def compare_model_weights_and_grads(self, base_model, model):
        for name, base_param in base_model.named_parameters():
            param = model.get_parameter(name)
            base_grad = base_param.grad.data
            grad = param.grad
            if isinstance(param, DTensor):
                param = param.redistribute(param.device_mesh, [Replicate()], async_op=False)._local_tensor

            torch.testing.assert_allclose(param, base_param)
            if isinstance(grad.data, DTensor):
                grad = grad.data.redistribute(grad.data.device_mesh, [Replicate()], async_op=False)._local_tensor
            torch.testing.assert_allclose(base_grad, grad, atol=1e-4, rtol=1e-4)
//  [python]
def compare_model_weights_and_grads(self, base_model, model):
        for name, base_param in base_model.named_parameters():
            param = model.get_parameter(name)
            base_grad = base_param.grad.data
            grad = param.grad
            if isinstance(param, DTensor):
                param = param.redistribute(param.device_mesh, [Replicate()], async_op=False)._local_tensor

            torch.testing.assert_close(param, base_param)
            if isinstance(grad.data, DTensor):
                grad = grad.data.redistribute(grad.data.device_mesh, [Replicate()], async_op=False)._local_tensor
            torch.testing.assert_close(base_grad, grad, atol=1e-4, rtol=1e-4)
//  [python]
def test_different_quant_descriptors_are_piped(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer
        from pytorch_quantization.calib import HistogramCalibrator
        from pytorch_quantization.calib import MaxCalibrator
        from pytorch_quantization.tensor_quant import QuantDescriptor

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        from pytorch_quantization.nn import QuantConv2d

        q_util.register_quantization_mapping(
            layer_names={"conv1"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="max"),
            weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
        )
        from pytorch_quantization.tensor_quant import QuantDescriptor

        q_util.register_quantization_mapping(
            layer_names={"conv2"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
            weights_quant_descriptor=QuantDescriptor(calib_method="max"),
        )
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._input_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_different_quant_descriptors_are_piped(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer
        from pytorch_quantization.calib import HistogramCalibrator
        from pytorch_quantization.calib import MaxCalibrator
        from pytorch_quantization.tensor_quant import QuantDescriptor

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        from pytorch_quantization.nn import QuantConv2d

        q_util.register_quantization_mapping(
            layer_names={"conv1"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="max"),
            weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
        )
        from pytorch_quantization.tensor_quant import QuantDescriptor

        q_util.register_quantization_mapping(
            layer_names={"conv2"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
            weights_quant_descriptor=QuantDescriptor(calib_method="max"),
        )
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._input_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_returns_expected_average(dataloader, expected_average):
        output_tensor = compute_average_features_from_images(dataloader, nn.Identity())
        torch.allclose(output_tensor, expected_average)
//  [python]
def test_returns_expected_average(dataloader, expected_average):
        output_tensor = compute_average_features_from_images(dataloader, nn.Identity())
        torch.testing.assert_close(output_tensor, expected_average)
//  [python]
def _test_wasserstein1d_with_input(
        self,
        compute_result: torch.Tensor,
        x: torch.Tensor,
        y: torch.Tensor,
        x_weights: Optional[torch.Tensor] = None,
        y_weights: Optional[torch.Tensor] = None,
    ) -> None:
        if x.ndim == 1:
            my_compute_result = wasserstein_1d(x, y, x_weights, y_weights)
            t.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)

                t.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
        else:
            my_compute_result = torch.tensor(
                [
                    wasserstein_1d(x, y, x_weights, y_weights)
                    for x, y, x_weights, y_weights in zip(x, y, x_weights, y_weights)
                ]
            ).to(x.device)
            t.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)

                t.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
//  [python]
def _test_wasserstein1d_with_input(
        self,
        compute_result: torch.Tensor,
        x: torch.Tensor,
        y: torch.Tensor,
        x_weights: Optional[torch.Tensor] = None,
        y_weights: Optional[torch.Tensor] = None,
    ) -> None:
        if x.ndim == 1:
            my_compute_result = wasserstein_1d(x, y, x_weights, y_weights)
            torch.testing.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)

                torch.testing.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
        else:
            my_compute_result = torch.tensor(
                [
                    wasserstein_1d(x, y, x_weights, y_weights)
                    for x, y, x_weights, y_weights in zip(x, y, x_weights, y_weights)
                ]
            ).to(x.device)
            torch.testing.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)

                torch.testing.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
//  [python]
def _test_sum_with_input(
        self,
        val: torch.Tensor,
        weight: Union[float, torch.Tensor] = 1.0,
    ) -> None:
        self.assertEqual(
            sum(val),
            torch.sum(val),
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_sum_with_input(
        self,
        val: torch.Tensor,
        weight: Union[float, torch.Tensor] = 1.0,
    ) -> None:
        torch.testing.assert_close(
            sum(val),
            torch.sum(val),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_predict_embeddings_returns_expected_dataframe(
        dataloader, expected_dataframe
    ):
        output_dataframe = predict_embeddings(dataloader, nn.Identity())
        pd.testing.assert_series_equal(
            output_dataframe.class_name, expected_dataframe.class_name
        )
        expected_dataframe_tensor = torch.stack(
            expected_dataframe.embedding.values.tolist()
        )
        output_dataframe_tensor = torch.stack(
            output_dataframe.embedding.values.tolist()
        )
        torch.allclose(output_dataframe_tensor, expected_dataframe_tensor)
//  [python]
def test_predict_embeddings_returns_expected_dataframe(
        dataloader, expected_dataframe
    ):
        output_dataframe = predict_embeddings(dataloader, nn.Identity())
        pd.testing.assert_series_equal(
            output_dataframe.class_name, expected_dataframe.class_name
        )
        expected_dataframe_tensor = torch.stack(
            expected_dataframe.embedding.values.tolist()
        )
        output_dataframe_tensor = torch.stack(
            output_dataframe.embedding.values.tolist()
        )
        torch.testing.assert_close(output_dataframe_tensor, expected_dataframe_tensor)
//  [python]
def test_encdec_multihead_attn_norm_add(self):
        grads = torch.randn_like(self.tst_inputs_q)

        for _ in range(5):
            ref_outputs, _ = self.ref_layer.forward(
                self.ref_inputs_q,
                self.ref_inputs_k,
                self.ref_inputs_k,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            tst_outputs, _ = self.tst_layer.forward(
                self.tst_inputs_q,
                self.tst_inputs_k,
                self.tst_inputs_k,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            self.ref_inputs_q.backward(grads)
            self.tst_inputs_q.backward(grads)

        torch.allclose(self.ref_inputs_q, self.tst_inputs_q, atol=1e-5, rtol=1e-5)
        torch.allclose(self.ref_inputs_k, self.tst_inputs_k, atol=1e-5, rtol=1e-5)
        torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_encdec_multihead_attn_norm_add(self):
        grads = torch.randn_like(self.tst_inputs_q)

        for _ in range(5):
            ref_outputs, _ = self.ref_layer.forward(
                self.ref_inputs_q,
                self.ref_inputs_k,
                self.ref_inputs_k,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            tst_outputs, _ = self.tst_layer.forward(
                self.tst_inputs_q,
                self.tst_inputs_k,
                self.tst_inputs_k,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            self.ref_inputs_q.backward(grads)
            self.tst_inputs_q.backward(grads)

        torch.testing.assert_close(self.ref_inputs_q, self.tst_inputs_q, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(self.ref_inputs_k, self.tst_inputs_k, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_sum_input_valid_weight(self) -> None:
        def _compute_result(
            val: torch.Tensor, weight: Union[float, torch.Tensor]
        ) -> torch.Tensor:
            weighted_sum = torch.tensor(0.0)
            if isinstance(weight, torch.Tensor):
                weight = weight.numpy().flatten()
            weighted_sum += val.numpy().flatten().dot(weight).sum()

            return weighted_sum

        inputs = [
            torch.rand(1),
            torch.rand(BATCH_SIZE, 4),
            torch.rand(BATCH_SIZE, 3, 4),
            torch.rand(5),
            torch.rand(10),
        ]
        weights = [
            torch.rand(1),
            torch.rand(BATCH_SIZE, 4),
            torch.rand(BATCH_SIZE, 3, 4),
            0.8,
            2,
        ]

        for input, weight in zip(inputs, weights):
            t.assert_allclose(
                sum(input, weight),
                _compute_result(input, weight),
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )
//  [python]
def test_sum_input_valid_weight(self) -> None:
        def _compute_result(
            val: torch.Tensor, weight: Union[float, torch.Tensor]
        ) -> torch.Tensor:
            weighted_sum = torch.tensor(0.0)
            if isinstance(weight, torch.Tensor):
                weight = weight.numpy().flatten()
            weighted_sum += val.numpy().flatten().dot(weight).sum()

            return weighted_sum

        inputs = [
            torch.rand(1),
            torch.rand(BATCH_SIZE, 4),
            torch.rand(BATCH_SIZE, 3, 4),
            torch.rand(5),
            torch.rand(10),
        ]
        weights = [
            torch.rand(1),
            torch.rand(BATCH_SIZE, 4),
            torch.rand(BATCH_SIZE, 3, 4),
            0.8,
            2,
        ]

        for input, weight in zip(inputs, weights):
            torch.testing.assert_close(
                sum(input, weight),
                _compute_result(input, weight),
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )
//  [python]
def _assert_abs_diff_small_base(compare_func, expected, actual, eps):
    # Helper function for comparing absolute differences of two torch tensors.
    abs_diff = torch.abs(expected - actual)
    err = compare_func(abs_diff)
    zero_tensor = torch.tensor(0, dtype=err.dtype)
    rtol = 1.6e-2 if err.dtype == torch.bfloat16 else 1.3e-6  
    torch.testing.assert_allclose(err, zero_tensor, atol=eps, rtol=rtol)
//  [python]
def _assert_abs_diff_small_base(compare_func, expected, actual, eps):
    # Helper function for comparing absolute differences of two torch tensors.
    abs_diff = torch.abs(expected - actual)
    err = compare_func(abs_diff)
    zero_tensor = torch.tensor(0, dtype=err.dtype)
    rtol = 1.6e-2 if err.dtype == torch.bfloat16 else 1.3e-6  
    torch.testing.assert_close(err, zero_tensor, atol=eps, rtol=rtol)
//  [python]
def initialize_and_test_parameters(
    rank: int,
    world_size: int,
    backend: str,
    embedding_tables: Union[EmbeddingCollection, EmbeddingBagCollection],
    sharding_type: str,
    sharders: List[ModuleSharder[nn.Module]],
    table_name: str,
    local_size: Optional[int] = None,
) -> None:
    with MultiProcessContext(rank, world_size, backend, local_size) as ctx:

        module_sharding_plan = construct_module_sharding_plan(
            embedding_tables,
            per_param_sharding={
                table_name: _select_sharding_type(sharding_type),
            },
            local_size=ctx.local_size,
            world_size=ctx.world_size,
            device_type=ctx.device.type,
        )

        model = DistributedModelParallel(
            module=embedding_tables,
            plan=ShardingPlan({"": module_sharding_plan}),
            # pyre-fixme[6]: For 1st argument expected `ProcessGroup` but got
            #  `Optional[ProcessGroup]`.
            env=ShardingEnv.from_process_group(ctx.pg),
            sharders=sharders,
            device=ctx.device,
        )

        key = (
            f"embeddings.{table_name}.weight"
            if isinstance(embedding_tables, EmbeddingCollection)
            else f"embedding_bags.{table_name}.weight"
        )

        if isinstance(model.state_dict()[key], ShardedTensor):
            if ctx.rank == 0:
                gathered_tensor = torch.empty_like(embedding_tables.state_dict()[key])
            else:
                gathered_tensor = None

            model.state_dict()[key].gather(dst=0, out=gathered_tensor)

            if ctx.rank == 0:
                torch.testing.assert_allclose(
                    gathered_tensor,
                    embedding_tables.state_dict()[key],
                )
        elif isinstance(model.state_dict()[key], torch.Tensor):
            torch.testing.assert_allclose(
                embedding_tables.state_dict()[key].cpu(),
                model.state_dict()[key].cpu(),
            )
        else:
            raise AssertionError(
                f"Model state dict contains unsupported type for key: {key}"
            )
//  [python]
def initialize_and_test_parameters(
    rank: int,
    world_size: int,
    backend: str,
    embedding_tables: Union[EmbeddingCollection, EmbeddingBagCollection],
    sharding_type: str,
    sharders: List[ModuleSharder[nn.Module]],
    table_name: str,
    local_size: Optional[int] = None,
) -> None:
    with MultiProcessContext(rank, world_size, backend, local_size) as ctx:

        module_sharding_plan = construct_module_sharding_plan(
            embedding_tables,
            per_param_sharding={
                table_name: _select_sharding_type(sharding_type),
            },
            local_size=ctx.local_size,
            world_size=ctx.world_size,
            device_type=ctx.device.type,
        )

        model = DistributedModelParallel(
            module=embedding_tables,
            plan=ShardingPlan({"": module_sharding_plan}),
            # pyre-fixme[6]: For 1st argument expected `ProcessGroup` but got
            #  `Optional[ProcessGroup]`.
            env=ShardingEnv.from_process_group(ctx.pg),
            sharders=sharders,
            device=ctx.device,
        )

        key = (
            f"embeddings.{table_name}.weight"
            if isinstance(embedding_tables, EmbeddingCollection)
            else f"embedding_bags.{table_name}.weight"
        )

        if isinstance(model.state_dict()[key], ShardedTensor):
            if ctx.rank == 0:
                gathered_tensor = torch.empty_like(embedding_tables.state_dict()[key])
            else:
                gathered_tensor = None

            model.state_dict()[key].gather(dst=0, out=gathered_tensor)

            if ctx.rank == 0:
                torch.testing.assert_close(
                    gathered_tensor,
                    embedding_tables.state_dict()[key],
                )
        elif isinstance(model.state_dict()[key], torch.Tensor):
            torch.testing.assert_close(
                embedding_tables.state_dict()[key].cpu(),
                model.state_dict()[key].cpu(),
            )
        else:
            raise AssertionError(
                f"Model state dict contains unsupported type for key: {key}"
            )
//  [python]
def test_nbit_ssd(self) -> None:
        import tempfile

        E = int(1e4)
        D = 128
        N = 100
        indices = torch.as_tensor(np.random.choice(E, replace=False, size=(N,)))
        weights = torch.empty(N, D, dtype=torch.uint8)
        output_weights = torch.empty_like(weights)
        count = torch.tensor([N])

        feature_table_map = list(range(1))
        emb = (
            ssd_split_table_batched_embeddings_ops.SSDIntNBitTableBatchedEmbeddingBags(
                embedding_specs=[("", E, D, SparseType.FP32)],
                feature_table_map=feature_table_map,
                ssd_storage_directory=tempfile.mkdtemp(),
                cache_sets=1,
            )
        )
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()

        emb.ssd_db.set_cuda(indices, weights, count, 1)
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        torch.testing.assert_allclose(weights, output_weights)
//  [python]
def test_nbit_ssd(self) -> None:
        import tempfile

        E = int(1e4)
        D = 128
        N = 100
        indices = torch.as_tensor(np.random.choice(E, replace=False, size=(N,)))
        weights = torch.empty(N, D, dtype=torch.uint8)
        output_weights = torch.empty_like(weights)
        count = torch.tensor([N])

        feature_table_map = list(range(1))
        emb = (
            ssd_split_table_batched_embeddings_ops.SSDIntNBitTableBatchedEmbeddingBags(
                embedding_specs=[("", E, D, SparseType.FP32)],
                feature_table_map=feature_table_map,
                ssd_storage_directory=tempfile.mkdtemp(),
                cache_sets=1,
            )
        )
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()

        emb.ssd_db.set_cuda(indices, weights, count, 1)
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        torch.testing.assert_close(weights, output_weights)
//  [python]
def test_ssd(self) -> None:
        import tempfile

        E = int(1e4)
        D = 128
        N = 100
        indices = torch.as_tensor(np.random.choice(E, replace=False, size=(N,)))
        weights = torch.randn(N, D)
        output_weights = torch.empty_like(weights)
        count = torch.tensor([N])

        feature_table_map = list(range(1))
        emb = ssd_split_table_batched_embeddings_ops.SSDTableBatchedEmbeddingBags(
            embedding_specs=[(E, D)],
            feature_table_map=feature_table_map,
            ssd_storage_directory=tempfile.mkdtemp(),
            cache_sets=1,
            ssd_uniform_init_lower=-0.1,
            ssd_uniform_init_upper=0.1,
        )
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        assert (output_weights <= 0.1).all().item()
        assert (output_weights >= -0.1).all().item()

        emb.ssd_db.set_cuda(indices, weights, count, 1)
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        torch.testing.assert_allclose(weights, output_weights)
//  [python]
def test_ssd(self) -> None:
        import tempfile

        E = int(1e4)
        D = 128
        N = 100
        indices = torch.as_tensor(np.random.choice(E, replace=False, size=(N,)))
        weights = torch.randn(N, D)
        output_weights = torch.empty_like(weights)
        count = torch.tensor([N])

        feature_table_map = list(range(1))
        emb = ssd_split_table_batched_embeddings_ops.SSDTableBatchedEmbeddingBags(
            embedding_specs=[(E, D)],
            feature_table_map=feature_table_map,
            ssd_storage_directory=tempfile.mkdtemp(),
            cache_sets=1,
            ssd_uniform_init_lower=-0.1,
            ssd_uniform_init_upper=0.1,
        )
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        assert (output_weights <= 0.1).all().item()
        assert (output_weights >= -0.1).all().item()

        emb.ssd_db.set_cuda(indices, weights, count, 1)
        emb.ssd_db.get_cuda(indices, output_weights, count)
        torch.cuda.synchronize()
        torch.testing.assert_close(weights, output_weights)
//  [python]
def check_jvp(f, *primals, executor, atol=None, rtol=None):
    """Check that the Jacobian-vector product of a function is correct.

    Args:
        f (callable): The function to differentiate.
        *primals (torch.Tensor): The input tensors.
        executor (str): The executor to use. Defaults to "torch".
        atol (float): Absolute tolerance. Defaults to None.
        rtol (float): Relative tolerance. Defaults to None.

    Raises:
        AssertionError: If the Jacobian-vector product is not correct.
    """
    tangents = tree_map(make_tensor_like, primals)
    actual_p, actual_t = executor.make_callable_legacy(jvp(f))(primals, tangents)
    expected_p, expected_t = numerical_jvp(executor.make_callable_legacy(f))(primals, tangents)
    t.assert_close(expected_p, actual_p, atol=atol, rtol=rtol)
    t.assert_close(expected_t, actual_t, atol=atol, rtol=rtol)
//  [python]
def check_jvp(f, *primals, executor, atol=None, rtol=None):
    """Check that the Jacobian-vector product of a function is correct.

    Args:
        f (callable): The function to differentiate.
        *primals (torch.Tensor): The input tensors.
        executor (str): The executor to use. Defaults to "torch".
        atol (float): Absolute tolerance. Defaults to None.
        rtol (float): Relative tolerance. Defaults to None.

    Raises:
        AssertionError: If the Jacobian-vector product is not correct.
    """
    tangents = tree_map(make_tensor_like, primals)
    actual_p, actual_t = executor.make_callable_legacy(jvp(f))(primals, tangents)
    expected_p, expected_t = numerical_jvp(executor.make_callable_legacy(f))(primals, tangents)
    torch.testing.assert_close(expected_p, actual_p, atol=atol, rtol=rtol)
    torch.testing.assert_close(expected_t, actual_t, atol=atol, rtol=rtol)
//  [python]
def test_lora_merge():
    config = Config(
        n_layer=1,
        n_head=2,
        n_embd=8,
        block_size=8,
        vocab_size=8,
        lora_r=8,
        lora_alpha=8,
        lora_dropout=0.1,
        lora_query=True,
        lora_value=True,
        lora_projection=True,
    )
    model = LoRAGPT(config)
    model.train()
    attn_proj = model.transformer.h[0].attn.proj

    initial_weight = attn_proj.linear.weight.clone()
    assert torch.equal(attn_proj.linear.weight, initial_weight)

    # perform an update to the LoRA weights
    mark_only_lora_as_trainable(model)
    optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
    y = model(torch.randint(0, 8, size=(2, 4), dtype=torch.int64))
    y.sum().backward()
    optimizer.step()
    optimizer.zero_grad()
    # the weight remains unchanged (only lora A and B change)
    assert torch.equal(attn_proj.linear.weight, initial_weight)

    # calling merge() multiple times in a row should not merge multiple times
    merge_lora_weights(model)
    assert attn_proj.merged
    weight_after = attn_proj.linear.weight.clone()
    merge_lora_weights(model)
    merge_lora_weights(model)
    assert torch.equal(attn_proj.linear.weight, weight_after)

    # check that `W_after = W_initial + (A x B)`
    delta_w = attn_proj.get_lora_AB()
    torch.testing.assert_allclose(weight_after, initial_weight + delta_w)
//  [python]
def test_lora_merge():
    config = Config(
        n_layer=1,
        n_head=2,
        n_embd=8,
        block_size=8,
        vocab_size=8,
        lora_r=8,
        lora_alpha=8,
        lora_dropout=0.1,
        lora_query=True,
        lora_value=True,
        lora_projection=True,
    )
    model = LoRAGPT(config)
    model.train()
    attn_proj = model.transformer.h[0].attn.proj

    initial_weight = attn_proj.linear.weight.clone()
    assert torch.equal(attn_proj.linear.weight, initial_weight)

    # perform an update to the LoRA weights
    mark_only_lora_as_trainable(model)
    optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
    y = model(torch.randint(0, 8, size=(2, 4), dtype=torch.int64))
    y.sum().backward()
    optimizer.step()
    optimizer.zero_grad()
    # the weight remains unchanged (only lora A and B change)
    assert torch.equal(attn_proj.linear.weight, initial_weight)

    # calling merge() multiple times in a row should not merge multiple times
    merge_lora_weights(model)
    assert attn_proj.merged
    weight_after = attn_proj.linear.weight.clone()
    merge_lora_weights(model)
    merge_lora_weights(model)
    assert torch.equal(attn_proj.linear.weight, weight_after)

    # check that `W_after = W_initial + (A x B)`
    delta_w = attn_proj.get_lora_AB()
    torch.testing.assert_close(weight_after, initial_weight + delta_w)
//  [python]
def test_convert_element_type_with_float(executor, device, _):
    # Verifies the fix for "grad transform hits error: AttributeError: 'float'
    # object has no attribute 'dtype'"
    from thunder.core.transforms import value_and_grad

    a = make_tensor([5], dtype=torch.float32, device=device)

    @value_and_grad
    def fn(t0):
        return t0 / 2

    out, (grad,) = executor.make_callable_legacy(fn)(a)
    t.assert_close(out, a / 2)
    t.assert_close(grad, torch.ones_like(a) / 2)
//  [python]
def test_convert_element_type_with_float(executor, device, _):
    # Verifies the fix for "grad transform hits error: AttributeError: 'float'
    # object has no attribute 'dtype'"
    from thunder.core.transforms import value_and_grad

    a = make_tensor([5], dtype=torch.float32, device=device)

    @value_and_grad
    def fn(t0):
        return t0 / 2

    out, (grad,) = executor.make_callable_legacy(fn)(a)
    torch.testing.assert_close(out, a / 2)
    torch.testing.assert_close(grad, torch.ones_like(a) / 2)
//  [python]
def test_torch_autograd_module(executor, device, _):
    l = torch.nn.Linear(3, 4, bias=False, device=device)
    a = make_tensor((2, 3), device=device, dtype=torch.float32, requires_grad=True)
    g = make_tensor((2, 4), device=device, dtype=torch.float32)

    for cache_mode in ("constant values", "same input"):
        lc = executor.make_callable(
            l,
            disable_torch_autograd=False,
            cache_mode=cache_mode,
        )
        lc.zero_grad()
        a.grad = None
        out = lc(a)
        out.backward(g)
        l_grad = l.weight.grad
        torch.testing.assert_allclose(l_grad, g.mT @ a)
        torch.testing.assert_allclose(a.grad, g @ l.weight)
//  [python]
def test_torch_autograd_module(executor, device, _):
    l = torch.nn.Linear(3, 4, bias=False, device=device)
    a = make_tensor((2, 3), device=device, dtype=torch.float32, requires_grad=True)
    g = make_tensor((2, 4), device=device, dtype=torch.float32)

    for cache_mode in ("constant values", "same input"):
        lc = executor.make_callable(
            l,
            disable_torch_autograd=False,
            cache_mode=cache_mode,
        )
        lc.zero_grad()
        a.grad = None
        out = lc(a)
        out.backward(g)
        l_grad = l.weight.grad
        torch.testing.assert_close(l_grad, g.mT @ a)
        torch.testing.assert_close(a.grad, g @ l.weight)
//  [python]
def test_torch_autograd_function_with_kwargs_static_caching(executor, device, _):
    def func(a, b):
        return a - b

    func = thunder.jit(func, executors=executor.executors_list(), disable_torch_autograd=False)

    a = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    b = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)

    # First call func(a, b) to populate the cache
    func(a, b)
    t.assert_close(func(a, b), a - b)
    t.assert_close(func(b=a, a=b), b - a)
    assert torch.autograd.gradcheck(lambda a, b: func(b=a, a=b), (a, b))
//  [python]
def test_torch_autograd_function_with_kwargs_static_caching(executor, device, _):
    def func(a, b):
        return a - b

    func = thunder.jit(func, executors=executor.executors_list(), disable_torch_autograd=False)

    a = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)
    b = make_tensor((2, 3), device=device, dtype=torch.float64, requires_grad=True)

    # First call func(a, b) to populate the cache
    func(a, b)
    torch.testing.assert_close(func(a, b), a - b)
    torch.testing.assert_close(func(b=a, a=b), b - a)
    assert torch.autograd.gradcheck(lambda a, b: func(b=a, a=b), (a, b))
//  [python]
def test_make_aug_forward_and_backward(executor, device, _):
    from thunder.core.vjp_utils import make_aug_forward_and_backward
    from thunder.core.prims import mul

    def fun(a, b):
        return mul(a, b)

    @executor.make_callable
    def expected_aug_fw(a, b):
        return fun(a, b), (a, b)

    @executor.make_callable
    def fun_bw(a, b, g):
        return {"a": g * b, "b": g * a}

    x = torch.tensor(2.0, device=device)
    y = torch.tensor(3.0, device=device)
    v = torch.tensor(1.5, device=device)

    trace = thunder.trace()(fun, x, y)
    mul_bsym = trace.bound_symbols[2]
    assert mul_bsym.sym.name == "mul"

    aug_fw, bw = make_aug_forward_and_backward(mul_bsym)
    aug_fw = executor.make_callable(aug_fw)
    actual_aug_fw, actual_saved = aug_fw(x, y)
    expected_aug_fw, expected_saved = expected_aug_fw(x, y)
    torch.testing.assert_allclose(actual_aug_fw, expected_aug_fw)

    bw = executor.make_callable(bw)
    actual_bw = bw(*actual_saved, v)
    expected_bw = fun_bw(*expected_saved, v)
    torch.testing.assert_allclose(actual_bw, expected_bw)
//  [python]
def test_make_aug_forward_and_backward(executor, device, _):
    from thunder.core.vjp_utils import make_aug_forward_and_backward
    from thunder.core.prims import mul

    def fun(a, b):
        return mul(a, b)

    @executor.make_callable
    def expected_aug_fw(a, b):
        return fun(a, b), (a, b)

    @executor.make_callable
    def fun_bw(a, b, g):
        return {"a": g * b, "b": g * a}

    x = torch.tensor(2.0, device=device)
    y = torch.tensor(3.0, device=device)
    v = torch.tensor(1.5, device=device)

    trace = thunder.trace()(fun, x, y)
    mul_bsym = trace.bound_symbols[2]
    assert mul_bsym.sym.name == "mul"

    aug_fw, bw = make_aug_forward_and_backward(mul_bsym)
    aug_fw = executor.make_callable(aug_fw)
    actual_aug_fw, actual_saved = aug_fw(x, y)
    expected_aug_fw, expected_saved = expected_aug_fw(x, y)
    torch.testing.assert_close(actual_aug_fw, expected_aug_fw)

    bw = executor.make_callable(bw)
    actual_bw = bw(*actual_saved, v)
    expected_bw = fun_bw(*expected_saved, v)
    torch.testing.assert_close(actual_bw, expected_bw)
//  [python]
def test_phantom_grad_unpack(executor, device: str, dtype: dtypes.dtype):
    # Tests tuple unpacking
    def foo(tup):
        a, b = tup
        return a * b

    cfoo = thunder.compile(foo, disable_preprocessing=True)
    cfoo_grad = grad(cfoo)

    a = torch.randn((2, 2), requires_grad=True)
    b = torch.randn((2, 2), requires_grad=True)

    a_grad, b_grad = cfoo_grad((a, b))

    t.assert_close(a_grad, b)
    t.assert_close(b_grad, a)

    # Tests dict unpacking
    def bar(d):
        a, b = d["a"], d["b"]
        return a * b

    cbar = thunder.compile(bar, disable_preprocessing=True)
    cbar_grad = grad(cbar)

    a_grad, b_grad = cbar_grad({"a": a, "b": b})

    t.assert_close(a_grad, b)
    t.assert_close(b_grad, a)
//  [python]
def test_phantom_grad_unpack(executor, device: str, dtype: dtypes.dtype):
    # Tests tuple unpacking
    def foo(tup):
        a, b = tup
        return a * b

    cfoo = thunder.compile(foo, disable_preprocessing=True)
    cfoo_grad = grad(cfoo)

    a = torch.randn((2, 2), requires_grad=True)
    b = torch.randn((2, 2), requires_grad=True)

    a_grad, b_grad = cfoo_grad((a, b))

    assert_close(a_grad, b)
    assert_close(b_grad, a)

    # Tests dict unpacking
    def bar(d):
        a, b = d["a"], d["b"]
        return a * b

    cbar = thunder.compile(bar, disable_preprocessing=True)
    cbar_grad = grad(cbar)

    a_grad, b_grad = cbar_grad({"a": a, "b": b})

    assert_close(a_grad, b)
    assert_close(b_grad, a)
//  [python]
def test_populate_grads_mlp(executor, device, dtype):
    from thunder.benchmarks import NanoGPTMLPBenchmark, NanoGPTConfig

    # NOTE Currently setting dropout to zero for reproducibility, other settings taken from gpt2 config
    config = NanoGPTConfig(dropout=0, n_layer=12, n_head=12, n_embd=768)

    bench = NanoGPTMLPBenchmark(config=config, requires_grad=True, device=device, dtype=dtype)
    model = bench.fn()
    (x,), kwargs = bench.make_batch()

    result = model(x)
    result.backward(torch.ones_like(result))
    torch_grads = extract_grads(model)

    clear_grads(model)

    tom = executor.make_callable(model)
    tom_grad = grad(tom)
    thunder_grads = tom_grad(x)

    populate_grads(thunder_grads, tom, args=(x,))
    thunder_grads = extract_grads(tom)

    t.assert_close(torch_grads, thunder_grads, atol=1e-3, rtol=1e-5)