jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_segmentation_model(model_fn, dev):
    set_rng_seed(0)
    defaults = {
        "num_classes": 10,
        "weights_backbone": None,
        "input_shape": (1, 3, 32, 32),
    }
    model_name = model_fn.__name__
    kwargs = {**defaults, **_model_params.get(model_name, {})}
    input_shape = kwargs.pop("input_shape")

    model = model_fn(**kwargs)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    with torch.no_grad(), freeze_rng_state():
        out = model(x)

    def check_out(out):
        prec = 0.01
        try:
            # We first try to assert the entire output if possible. This is not
            # only the best way to assert results but also handles the cases
            # where we need to create a new expected result.
            _assert_expected(out.cpu(), model_name, prec=prec)
        except AssertionError:
            # Unfortunately some segmentation models are flaky with autocast
            # so instead of validating the probability scores, check that the class
            # predictions match.
            expected_file = _get_expected_file(model_name)
            expected = torch.load(expected_file, weights_only=True)
            torch.testing.assert_allclose(
                out.argmax(dim=1), expected.argmax(dim=1), rtol=prec, atol=prec, check_device=False
            )
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out["out"])

    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)

    if dev == "cuda":
        with torch.cuda.amp.autocast(), torch.no_grad(), freeze_rng_state():
            out = model(x)
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out["out"])

    if not full_validation:
        msg = (
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
            "significant changes to the codebase."
        )
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)

    _check_input_backprop(model, x)
//  [python]
def test_segmentation_model(model_fn, dev):
    set_rng_seed(0)
    defaults = {
        "num_classes": 10,
        "weights_backbone": None,
        "input_shape": (1, 3, 32, 32),
    }
    model_name = model_fn.__name__
    kwargs = {**defaults, **_model_params.get(model_name, {})}
    input_shape = kwargs.pop("input_shape")

    model = model_fn(**kwargs)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    with torch.no_grad(), freeze_rng_state():
        out = model(x)

    def check_out(out):
        prec = 0.01
        try:
            # We first try to assert the entire output if possible. This is not
            # only the best way to assert results but also handles the cases
            # where we need to create a new expected result.
            _assert_expected(out.cpu(), model_name, prec=prec)
        except AssertionError:
            # Unfortunately some segmentation models are flaky with autocast
            # so instead of validating the probability scores, check that the class
            # predictions match.
            expected_file = _get_expected_file(model_name)
            expected = torch.load(expected_file, weights_only=True)
            torch.testing.assert_close(
                out.argmax(dim=1), expected.argmax(dim=1), rtol=prec, atol=prec, check_device=False
            )
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out["out"])

    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)

    if dev == "cuda":
        with torch.cuda.amp.autocast(), torch.no_grad(), freeze_rng_state():
            out = model(x)
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out["out"])

    if not full_validation:
        msg = (
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
            "significant changes to the codebase."
        )
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)

    _check_input_backprop(model, x)
//  [python]
def test__int_mm(self):
        # TODO(future): figure out what here needs to move to PT core,
        # if it's not already tested there

        m, k, n = 32, 32, 16
        x = torch.randint(-128, 127, (m, k), dtype=torch.int8, device="cuda")
        w = torch.randint(-128, 127, (k, n), dtype=torch.int8, device="cuda")

        y_ref = torch.matmul(x.float(), w.float()).to(torch.int32)
        y_raw = safe_int_mm(x, w)

        wrap_in_mm_opt = torch.compile(safe_int_mm, mode="max-autotune")
        # note: triton chokes on the line below on k == 8 and n == 8 with
        # https://www.internalfb.com/phabricator/paste/view/P683467944
        # TODO(future): file an issue
        y_opt = wrap_in_mm_opt(x, w)

        t.assert_close(y_ref, y_raw, atol=0, rtol=0)
        t.assert_close(y_ref, y_opt, atol=0, rtol=0)
//  [python]
def test__int_mm(self):
        # TODO(future): figure out what here needs to move to PT core,
        # if it's not already tested there

        m, k, n = 32, 32, 16
        x = torch.randint(-128, 127, (m, k), dtype=torch.int8, device="cuda")
        w = torch.randint(-128, 127, (k, n), dtype=torch.int8, device="cuda")

        y_ref = torch.matmul(x.float(), w.float()).to(torch.int32)
        y_raw = safe_int_mm(x, w)

        wrap_in_mm_opt = torch.compile(safe_int_mm, mode="max-autotune")
        # note: triton chokes on the line below on k == 8 and n == 8 with
        # https://www.internalfb.com/phabricator/paste/view/P683467944
        # TODO(future): file an issue
        y_opt = wrap_in_mm_opt(x, w)

        torch.testing.assert_close(y_ref, y_raw, atol=0, rtol=0)
        torch.testing.assert_close(y_ref, y_opt, atol=0, rtol=0)
//  [python]
def test_encdec_multihead_attn(self):
        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,
        )
        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)

        with torch.no_grad():
            ref_grads = torch.randn_like(ref_outputs)
            tst_grads = ref_grads.clone()
        ref_outputs.backward(ref_grads)
        tst_outputs.backward(tst_grads)
        torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_encdec_multihead_attn(self):
        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,
        )
        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)

        with torch.no_grad():
            ref_grads = torch.randn_like(ref_outputs)
            tst_grads = ref_grads.clone()
        ref_outputs.backward(ref_grads)
        tst_outputs.backward(tst_grads)
        torch.testing.assert_close(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)
//  [python]
def _test_dynamic_quant_per_tensor_numerics_impl(
        self, qmin, qmax, int_dtype, qint_dtype, float_dtype, device, qscheme
    ):
        x = torch.randn(256, dtype=float_dtype, device=device)
        y_vals, y_scale, y_zero_point = dynamically_quantize_per_tensor(
            x, qmin, qmax, int_dtype, qscheme
        )

        # reference
        # quantize_per_tensor_dynamic doesn't work for half, so we cast there and back
        x_for_ref = x.half().float() if float_dtype == torch.float16 else x

        # quantize_per_tensor_dynamic doesn't support qscheme, so we just do dynamic
        # quant manually with observers + static quant
        obs = MinMaxObserver(
            dtype=qint_dtype, qscheme=qscheme, quant_min=qmin, quant_max=qmax
        ).to(device)
        obs(x_for_ref)
        ref_scale, ref_zero_point = obs.calculate_qparams()
        y_ref = torch.quantize_per_tensor(
            x_for_ref, ref_scale, ref_zero_point, qint_dtype
        )

        # y_ref = torch.quantize_per_tensor_dynamic(x_for_ref, qint_dtype, False)
        # print(y_ref)
        if float_dtype == torch.float:
            assert torch.equal(y_vals, y_ref.int_repr())
        else:
            # numerics are not exactly aligned yet, off-by-one probably due
            # to rounding
            assert torch.max(torch.abs(y_vals - y_ref.int_repr())).item() <= 1
        self.assertEqual(
            y_scale, torch.tensor([y_ref.q_scale()], device=device, dtype=float_dtype)
        )
        if y_zero_point is not None:
            self.assertEqual(
                y_zero_point, torch.tensor([y_ref.q_zero_point()], device=device)
            )
        else:
            self.assertTrue(y_ref.q_zero_point() == 0)

        # dequantize and check again
        x_dq = dequantize_per_tensor(y_vals, y_scale, y_zero_point, float_dtype)
        y_ref_dq = y_ref.dequantize().to(float_dtype)
        if float_dtype == torch.float:
            self.assertEqual(x_dq, y_ref_dq)
        else:
            sqnr = compute_error(x_dq, y_ref_dq)
            self.assertTrue(sqnr.item() > 45.0)
//  [python]
def _test_dynamic_quant_per_tensor_numerics_impl(
        self, qmin, qmax, int_dtype, qint_dtype, float_dtype, device, qscheme
    ):
        x = torch.randn(256, dtype=float_dtype, device=device)
        y_vals, y_scale, y_zero_point = dynamically_quantize_per_tensor(
            x, qmin, qmax, int_dtype, qscheme
        )

        # reference
        # quantize_per_tensor_dynamic doesn't work for half, so we cast there and back
        x_for_ref = x.half().float() if float_dtype == torch.float16 else x

        # quantize_per_tensor_dynamic doesn't support qscheme, so we just do dynamic
        # quant manually with observers + static quant
        obs = MinMaxObserver(
            dtype=qint_dtype, qscheme=qscheme, quant_min=qmin, quant_max=qmax
        ).to(device)
        obs(x_for_ref)
        ref_scale, ref_zero_point = obs.calculate_qparams()
        y_ref = torch.quantize_per_tensor(
            x_for_ref, ref_scale, ref_zero_point, qint_dtype
        )

        # y_ref = torch.quantize_per_tensor_dynamic(x_for_ref, qint_dtype, False)
        # print(y_ref)
        if float_dtype == torch.float:
            assert torch.equal(y_vals, y_ref.int_repr())
        else:
            # numerics are not exactly aligned yet, off-by-one probably due
            # to rounding
            assert torch.max(torch.abs(y_vals - y_ref.int_repr())).item() <= 1
        torch.testing.assert_close(
            y_scale, torch.tensor([y_ref.q_scale()], device=device, dtype=float_dtype)
        )
        if y_zero_point is not None:
            assert torch.equal(
                y_zero_point, torch.tensor([y_ref.q_zero_point()], device=device)
            )
        else:
            self.assertTrue(y_ref.q_zero_point() == 0)

        # dequantize and check again
        x_dq = dequantize_per_tensor(y_vals, y_scale, y_zero_point, float_dtype)
        y_ref_dq = y_ref.dequantize().to(float_dtype)
        if float_dtype == torch.float:
            torch.testing.assert_close(x_dq, y_ref_dq)
        else:
            sqnr = compute_error(x_dq, y_ref_dq)
            self.assertTrue(sqnr.item() > 45.0)
//  [python]
def test_weight_t_and_non_t_numerics_match(self):
        # verify that numerics match whether weight is stored
        # in transposed format (for cuBLAS) vs non-transposed format
        # (for torch.compile)
        dtype = torch.half
        device = "cuda"
        lin_ref = nn.Linear(32, 16, dtype=dtype, device=device)
        lin_eager_t = copy.deepcopy(lin_ref)
        lin_opt_t = copy.deepcopy(lin_eager_t)
        lin_opt = copy.deepcopy(lin_eager_t)
        lin_eager_t = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_eager_t)
        lin_opt_t = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_opt_t)
        lin_opt = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_opt)
        lin_opt.store_w_int_repr_t = False

        x = torch.randn(32, 32, dtype=dtype, device=device)

        y_calib_eager_t = lin_eager_t(x)
        y_calib_opt_t = lin_opt_t(x)
        y_calib_opt = lin_opt(x)
        torch.testing.assert_allclose(y_calib_eager_t, y_calib_opt_t)
        torch.testing.assert_allclose(y_calib_eager_t, y_calib_opt)

        lin_eager_t.to_inference()
        lin_opt_t.to_inference()
        lin_opt.to_inference()

        torch.testing.assert_allclose(lin_eager_t.W_int_repr, lin_opt_t.W_int_repr)
        torch.testing.assert_allclose(lin_eager_t.W_int_repr, lin_opt.W_int_repr)

        lin_opt_t = torch.compile(lin_opt_t, mode="max-autotune")
        lin_opt = torch.compile(lin_opt, mode="max-autotune")

        y_ref = lin_ref(x)
        y_eager = lin_eager_t(x)
        y_opt_t = lin_opt_t(x)
        y_opt = lin_opt(x)

        if not torch.any(torch.isinf(y_ref)) and torch.any(torch.isinf(y_eager)):
            # eager mode torch._int_mm is sometimes buggy, when this happens
            # we can't really compare the compiled version against it properly
            print("eager mode torch._int_mm known bad, test is inconclusive")
            return

        sqnr_ref_eager = compute_error(y_ref, y_eager)
        sqnr_eager_opt_t = compute_error(y_eager, y_opt_t)
        sqnr_eager_opt = compute_error(y_eager, y_opt)
        # since torch.compile for a torch.half model can
        # change numerics significantly, we can only test for a high SQNR here
        # and not for closeness
        self.assertTrue(sqnr_eager_opt_t >= 45.0)
        self.assertTrue(sqnr_eager_opt >= 45.0)
        # y_opt_t and y_opt should be equivalent
        torch.testing.assert_allclose(y_opt_t, y_opt)
//  [python]
def test_weight_t_and_non_t_numerics_match(self):
        # verify that numerics match whether weight is stored
        # in transposed format (for cuBLAS) vs non-transposed format
        # (for torch.compile)
        dtype = torch.half
        device = "cuda"
        lin_ref = nn.Linear(32, 16, dtype=dtype, device=device)
        lin_eager_t = copy.deepcopy(lin_ref)
        lin_opt_t = copy.deepcopy(lin_eager_t)
        lin_opt = copy.deepcopy(lin_eager_t)
        lin_eager_t = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_eager_t)
        lin_opt_t = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_opt_t)
        lin_opt = SmoothFakeDynamicallyQuantizedLinear.from_float(lin_opt)
        lin_opt.store_w_int_repr_t = False

        x = torch.randn(32, 32, dtype=dtype, device=device)

        y_calib_eager_t = lin_eager_t(x)
        y_calib_opt_t = lin_opt_t(x)
        y_calib_opt = lin_opt(x)
        torch.testing.assert_close(y_calib_eager_t, y_calib_opt_t)
        torch.testing.assert_close(y_calib_eager_t, y_calib_opt)

        lin_eager_t.to_inference()
        lin_opt_t.to_inference()
        lin_opt.to_inference()

        torch.testing.assert_close(lin_eager_t.W_int_repr, lin_opt_t.W_int_repr)
        torch.testing.assert_close(lin_eager_t.W_int_repr, lin_opt.W_int_repr)

        lin_opt_t = torch.compile(lin_opt_t, mode="max-autotune")
        lin_opt = torch.compile(lin_opt, mode="max-autotune")

        y_ref = lin_ref(x)
        y_eager = lin_eager_t(x)
        y_opt_t = lin_opt_t(x)
        y_opt = lin_opt(x)

        if not torch.any(torch.isinf(y_ref)) and torch.any(torch.isinf(y_eager)):
            # eager mode torch._int_mm is sometimes buggy, when this happens
            # we can't really compare the compiled version against it properly
            print("eager mode torch._int_mm known bad, test is inconclusive")
            return

        sqnr_ref_eager = compute_error(y_ref, y_eager)
        sqnr_eager_opt_t = compute_error(y_eager, y_opt_t)
        sqnr_eager_opt = compute_error(y_eager, y_opt)
        # since torch.compile for a torch.half model can
        # change numerics significantly, we can only test for a high SQNR here
        # and not for closeness
        self.assertTrue(sqnr_eager_opt_t >= 45.0)
        self.assertTrue(sqnr_eager_opt >= 45.0)
        # y_opt_t and y_opt should be equivalent
        torch.testing.assert_close(y_opt_t, y_opt)
//  [python]
def _test_dynamic_quant_per_channel_numerics_impl(
        self, qmin, qmax, int_dtype, qint_dtype, float_dtype, device
    ):
        # verifies that dynamic quant per channel in plain pytorch matches
        # numerics of production AO code
        # TODO(future): test this on cpu-half, need to first make
        # torch.aminmax support half on cpu

        x = torch.randn(16, 32, device=device, dtype=float_dtype)
        y_vals, y_scale, y_zero_point = dynamically_quantize_per_channel(
            x, qmin, qmax, int_dtype
        )

        min_val, max_val = torch.aminmax(x, dim=1)

        # reference
        weight_obs = torch.ao.quantization.MovingAveragePerChannelMinMaxObserver(
            dtype=qint_dtype,
            quant_min=qmin,
            quant_max=qmax,
            qscheme=torch.per_channel_symmetric,
            averaging_constant=1.0,  # make it ignore previous iterations
        )
        weight_obs(x)
        y_ref_scale, y_ref_zp = weight_obs.calculate_qparams()
        y_ref_scale = y_ref_scale.to(device)
        y_ref_zp = y_ref_zp.to(device)
        # quantize_per_channel doesn't work for half, so we cast there and back
        x_for_ref = x.half().float() if float_dtype == torch.float16 else x
        y_ref = torch.quantize_per_channel(
            x_for_ref, y_ref_scale, y_ref_zp, 0, qint_dtype
        )

        torch.testing.assert_allclose(
            y_scale, y_ref.q_per_channel_scales().to(float_dtype)
        )
        assert torch.equal(y_zero_point, y_ref.q_per_channel_zero_points())
        # this test case has one element where the rounding is off by one
        # from Python-only code vs the c++ code, it's easy to repro with
        # various shapes.
        # Discussion here is relevant: https://github.com/pytorch/pytorch/issues/16498
        # TODO(future): figure out what to do about this
        # assert torch.equal(int_vals, q_reference.int_repr())
        assert torch.max(torch.abs(y_vals - y_ref.int_repr())) <= 1

        # dequantize
        x_dq = dequantize_per_channel(y_vals, y_scale, y_zero_point)
        x_ref_dq = y_ref.dequantize()
        # off-by-one for scale is okay
        torch.testing.assert_allclose(
            x_dq, x_ref_dq, atol=torch.max(y_scale).item() * 1.01, rtol=0.0001
        )
//  [python]
def _test_dynamic_quant_per_channel_numerics_impl(
        self, qmin, qmax, int_dtype, qint_dtype, float_dtype, device
    ):
        # verifies that dynamic quant per channel in plain pytorch matches
        # numerics of production AO code
        # TODO(future): test this on cpu-half, need to first make
        # torch.aminmax support half on cpu

        x = torch.randn(16, 32, device=device, dtype=float_dtype)
        y_vals, y_scale, y_zero_point = dynamically_quantize_per_channel(
            x, qmin, qmax, int_dtype
        )

        min_val, max_val = torch.aminmax(x, dim=1)

        # reference
        weight_obs = torch.ao.quantization.MovingAveragePerChannelMinMaxObserver(
            dtype=qint_dtype,
            quant_min=qmin,
            quant_max=qmax,
            qscheme=torch.per_channel_symmetric,
            averaging_constant=1.0,  # make it ignore previous iterations
        )
        weight_obs(x)
        y_ref_scale, y_ref_zp = weight_obs.calculate_qparams()
        y_ref_scale = y_ref_scale.to(device)
        y_ref_zp = y_ref_zp.to(device)
        # quantize_per_channel doesn't work for half, so we cast there and back
        x_for_ref = x.half().float() if float_dtype == torch.float16 else x
        y_ref = torch.quantize_per_channel(
            x_for_ref, y_ref_scale, y_ref_zp, 0, qint_dtype
        )

        torch.testing.assert_close(
            y_scale, y_ref.q_per_channel_scales().to(float_dtype)
        )
        assert torch.equal(y_zero_point, y_ref.q_per_channel_zero_points())
        # this test case has one element where the rounding is off by one
        # from Python-only code vs the c++ code, it's easy to repro with
        # various shapes.
        # Discussion here is relevant: https://github.com/pytorch/pytorch/issues/16498
        # TODO(future): figure out what to do about this
        # assert torch.equal(int_vals, q_reference.int_repr())
        assert torch.max(torch.abs(y_vals - y_ref.int_repr())) <= 1

        # dequantize
        x_dq = dequantize_per_channel(y_vals, y_scale, y_zero_point)
        x_ref_dq = y_ref.dequantize()
        # off-by-one for scale is okay
        torch.testing.assert_close(
            x_dq, x_ref_dq, atol=torch.max(y_scale).item() * 1.01, rtol=0.0001
        )
//  [python]
def test_encdec_multihead_attn_time_mask(self):
        grads = torch.randn_like(self.tst_inputs_q)
        time_mask_byte = torch.triu(
            torch.ones(
                self.tst_inputs_q.size(0), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        time_mask_bool = time_mask_byte.to(torch.bool)

        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=time_mask_bool,
            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=time_mask_byte,
            is_training=True,
        )

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

        torch.testing.assert_allclose(self.ref_inputs_q, self.tst_inputs_q, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(self.ref_inputs_k, self.tst_inputs_k, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_encdec_multihead_attn_time_mask(self):
        grads = torch.randn_like(self.tst_inputs_q)
        time_mask_byte = torch.triu(
            torch.ones(
                self.tst_inputs_q.size(0), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        time_mask_bool = time_mask_byte.to(torch.bool)

        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=time_mask_bool,
            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=time_mask_byte,
            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_train_load():
    tmpdir = "tests/tmp_load"
    solver = SupervisedSolver(problem=problem,
                model=model,
                extra_features=None,
                loss=LpLoss())
    trainer = Trainer(solver=solver,
                      max_epochs=15,
                      accelerator='cpu',
                      default_root_dir=tmpdir)
    trainer.train()
    new_solver = SupervisedSolver.load_from_checkpoint(
        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
        problem = problem, model=model)
    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
    assert new_solver.forward(test_pts).shape == (20, 1)
    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
    torch.testing.assert_allclose(
        new_solver.forward(test_pts),
        solver.forward(test_pts))
    import shutil
    shutil.rmtree(tmpdir)
//  [python]
def test_train_load():
    tmpdir = "tests/tmp_load"
    solver = SupervisedSolver(problem=problem,
                model=model,
                extra_features=None,
                loss=LpLoss())
    trainer = Trainer(solver=solver,
                      max_epochs=15,
                      accelerator='cpu',
                      default_root_dir=tmpdir)
    trainer.train()
    new_solver = SupervisedSolver.load_from_checkpoint(
        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
        problem = problem, model=model)
    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
    assert new_solver.forward(test_pts).shape == (20, 1)
    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
    torch.testing.assert_close(
        new_solver.forward(test_pts),
        solver.forward(test_pts))
    import shutil
    shutil.rmtree(tmpdir)
//  [python]
def test_encdec_multihead_attn_pad_mask(self):
        grads = torch.randn_like(self.tst_inputs_q)
        pad_mask_byte = torch.tril(
            torch.ones(
                self.tst_inputs_k.size(1), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        pad_mask_bool = pad_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs_q,
            self.ref_inputs_k,
            self.ref_inputs_k,
            key_padding_mask=pad_mask_bool,
            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=pad_mask_byte,
            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_pad_mask(self):
        grads = torch.randn_like(self.tst_inputs_q)
        pad_mask_byte = torch.tril(
            torch.ones(
                self.tst_inputs_k.size(1), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        pad_mask_bool = pad_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs_q,
            self.ref_inputs_k,
            self.ref_inputs_k,
            key_padding_mask=pad_mask_bool,
            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=pad_mask_byte,
            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_paillier_torch():
    import torch  # noqa: F401

    from aijack.defense.paillier import (  # noqa: F401
        PaillierKeyGenerator,
        PaillierTensor,
    )

    keygenerator = PaillierKeyGenerator(512)
    pk, sk = keygenerator.generate_keypair()

    ct_1 = pk.encrypt(13)
    ct_2 = pk.encrypt(0.5)
    ct_3 = ct_1 + ct_2

    pt_1 = PaillierTensor([ct_1, ct_2, ct_3])
    torch.testing.assert_allclose(
        pt_1.decrypt(sk), torch.Tensor([13, 0.5, 13.5]), atol=1e-5, rtol=1
    )

    pt_2 = pt_1 + torch.Tensor([0.4, 0.1, 0.2])
    torch.testing.assert_allclose(
        pt_2.decrypt(sk), torch.Tensor([13.4, 0.6, 13.7]), atol=1e-5, rtol=1
    )

    pt_3 = pt_1 * torch.Tensor([1, 2.5, 0.5])
    torch.testing.assert_allclose(
        pt_3.decrypt(sk), torch.Tensor([13, 1.25, 6.75]), atol=1e-5, rtol=1
    )

    pt_4 = pt_1 - torch.Tensor([0.7, 0.3, 0.6])
    torch.testing.assert_allclose(
        pt_4.decrypt(sk), torch.Tensor([14.3, 0.2, 12.9]), atol=1e-5, rtol=1
    )

    pt_5 = pt_1 * 2
    torch.testing.assert_allclose(
        pt_5.decrypt(sk), torch.Tensor([26, 1, 27]), atol=1e-5, rtol=1
    )
//  [python]
def test_paillier_torch():
    import torch  # noqa: F401

    from aijack.defense.paillier import (  # noqa: F401
        PaillierKeyGenerator,
        PaillierTensor,
    )

    keygenerator = PaillierKeyGenerator(512)
    pk, sk = keygenerator.generate_keypair()

    ct_1 = pk.encrypt(13)
    ct_2 = pk.encrypt(0.5)
    ct_3 = ct_1 + ct_2

    pt_1 = PaillierTensor([ct_1, ct_2, ct_3])
    torch.testing.assert_close(
        pt_1.decrypt(sk), torch.Tensor([13, 0.5, 13.5]), atol=1e-5, rtol=1
    )

    pt_2 = pt_1 + torch.Tensor([0.4, 0.1, 0.2])
    torch.testing.assert_close(
        pt_2.decrypt(sk), torch.Tensor([13.4, 0.6, 13.7]), atol=1e-5, rtol=1
    )

    pt_3 = pt_1 * torch.Tensor([1, 2.5, 0.5])
    torch.testing.assert_close(
        pt_3.decrypt(sk), torch.Tensor([13, 1.25, 6.75]), atol=1e-5, rtol=1
    )

    pt_4 = pt_1 - torch.Tensor([0.7, 0.3, 0.6])
    torch.testing.assert_close(
        pt_4.decrypt(sk), torch.Tensor([14.3, 0.2, 12.9]), atol=1e-5, rtol=1
    )

    pt_5 = pt_1 * 2
    torch.testing.assert_close(
        pt_5.decrypt(sk), torch.Tensor([26, 1, 27]), atol=1e-5, rtol=1
    )
//  [python]
def get_multi_weights_col(self, prefixes: List[Union[str, Tuple]], quantize: str, dim: int):
        if quantize in ["gptq", "awq"]:
            try:
                qweight = torch.cat(self.get_sharded_list("qweight", prefixes, dim=1), dim=1)
            except RuntimeError:
                raise RuntimeError(
                    "Cannot load `gptq` weight, make sure the model is already quantized, or quantize it with `lorax-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
                )

            qzeros = torch.cat(self.get_sharded_list("qzeros", prefixes, dim=1), dim=1)
            scales = torch.cat(self.get_sharded_list("scales", prefixes, dim=1), dim=1)
            if quantize == "gptq":
                # no tensor parallelism, so remove the range if provided
                prefixes = [p[0] if isinstance(p, tuple) else p for p in prefixes]
                w = [self.get_tensor(f"{p}.g_idx") for p in prefixes]
                for w2 in w[1:]:
                    torch.allclose(w2, w[0])
                g_idx = w[0]
            else:
                g_idx = None

            bits, groupsize = self._get_bits_and_groupsize()
            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, False)
        else:
            w = self.get_sharded_list("weight", prefixes, dim=0)
            weight = torch.cat(w, dim=dim)
        return weight
//  [python]
def get_multi_weights_col(self, prefixes: List[Union[str, Tuple]], quantize: str, dim: int):
        if quantize in ["gptq", "awq"]:
            try:
                qweight = torch.cat(self.get_sharded_list("qweight", prefixes, dim=1), dim=1)
            except RuntimeError:
                raise RuntimeError(
                    "Cannot load `gptq` weight, make sure the model is already quantized, or quantize it with `lorax-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
                )

            qzeros = torch.cat(self.get_sharded_list("qzeros", prefixes, dim=1), dim=1)
            scales = torch.cat(self.get_sharded_list("scales", prefixes, dim=1), dim=1)
            if quantize == "gptq":
                # no tensor parallelism, so remove the range if provided
                prefixes = [p[0] if isinstance(p, tuple) else p for p in prefixes]
                w = [self.get_tensor(f"{p}.g_idx") for p in prefixes]
                for w2 in w[1:]:
                    torch.testing.assert_close(w2, w[0])
                g_idx = w[0]
            else:
                g_idx = None

            bits, groupsize = self._get_bits_and_groupsize()
            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, False)
        else:
            w = self.get_sharded_list("weight", prefixes, dim=0)
            weight = torch.cat(w, dim=dim)
        return weight
//  [python]
def validate_cutmix(x, y, indices, x_cutmix, y_perm, bbox):
    # Create shuffled version of x, y for reference checking
    x_perm = x[indices]
    y_perm_ref = y[indices]

    # Explicitly check that the pixels and labels have been mixed correctly.
    for i in range(x.size(0)):  # Grab N
        # Check every pixel of the input data
        for j in range(x.size(2)):
            for k in range(x.size(3)):
                if (j >= bbox[0] and j < bbox[2]) and (k >= bbox[1] and k < bbox[3]):
                    t.testing.assert_close(x_perm[i, :, j, k], x_cutmix[i, :, j, k])
                else:
                    t.testing.assert_close(x[i, :, j, k], x_cutmix[i, :, j, k])
    # Check the label
    t.testing.assert_close(y_perm_ref, y_perm)
//  [python]
def validate_cutmix(x, y, indices, x_cutmix, y_perm, bbox):
    # Create shuffled version of x, y for reference checking
    x_perm = x[indices]
    y_perm_ref = y[indices]

    # Explicitly check that the pixels and labels have been mixed correctly.
    for i in range(x.size(0)):  # Grab N
        # Check every pixel of the input data
        for j in range(x.size(2)):
            for k in range(x.size(3)):
                if (j >= bbox[0] and j < bbox[2]) and (k >= bbox[1] and k < bbox[3]):
                    torch.testing.assert_close(x_perm[i, :, j, k], x_cutmix[i, :, j, k])
                else:
                    torch.testing.assert_close(x[i, :, j, k], x_cutmix[i, :, j, k])
    # Check the label
    torch.testing.assert_close(y_perm_ref, y_perm)
//  [python]
def test_distributed_state_dict(self) -> None:
        state_dict_with_param_groups = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )
        self.assertEqual(state_dict_with_param_groups.keys(), {"state", "param_groups"})

        self.assertEqual(state_dict_with_param_groups["state"], self._distributed_state_dict["state"])
        self.assertEqual(
            state_dict_with_param_groups["param_groups"],
            self._distributed_state_dict["param_groups"],
        )

        state_dict_without_param_groups = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=False,
        )
        self.assertEqual(state_dict_without_param_groups.keys(), {"state"})

        self.assertEqual(
            state_dict_without_param_groups["state"],
            self._distributed_state_dict["state"],
        )
//  [python]
def test_distributed_state_dict(self) -> None:
        state_dict_with_param_groups = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )
        self.assertEqual(state_dict_with_param_groups.keys(), {"state", "param_groups"})

        torch.testing.assert_close(
            state_dict_with_param_groups["state"], self._distributed_state_dict["state"]
        )
        self.assertEqual(
            state_dict_with_param_groups["param_groups"],
            self._distributed_state_dict["param_groups"],
        )

        state_dict_without_param_groups = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=False,
        )
        self.assertEqual(state_dict_without_param_groups.keys(), {"state"})

        torch.testing.assert_close(
            state_dict_without_param_groups["state"],
            self._distributed_state_dict["state"],
        )
//  [python]
def test_load_distributed_state_dict(self) -> None:
        expected_distributed_state_dict = copy.deepcopy(self._distributed_state_dict)

        self._optimizer.load_distributed_state_dict(
            state_dict=self._distributed_state_dict,
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )

        actual_state_dict = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )

        self.assertEqual(
            actual_state_dict.keys(), expected_distributed_state_dict.keys()
        )
        self.assertEqual(
            actual_state_dict["param_groups"],
            expected_distributed_state_dict["param_groups"],
        ) 
        t.assert_allclose(
            actual_state_dict["state"], expected_distributed_state_dict["state"]
        )
//  [python]
def test_load_distributed_state_dict(self) -> None:
        expected_distributed_state_dict = copy.deepcopy(self._distributed_state_dict)

        self._optimizer.load_distributed_state_dict(
            state_dict=self._distributed_state_dict,
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )

        actual_state_dict = self._optimizer.distributed_state_dict(
            key_to_param=self._model.named_parameters(),
            save_param_groups=True,
        )

        self.assertEqual(
            actual_state_dict.keys(), expected_distributed_state_dict.keys()
        )
        torch.testing.assert_close(
            actual_state_dict["state"], expected_distributed_state_dict["state"]
        )
        self.assertEqual(
            actual_state_dict["param_groups"],
            expected_distributed_state_dict["param_groups"],
        )
//  [python]
def test_basic(self):
        """
        Check that we convert a sequence of reshape(unsqueeze(...)) into a
        single reshape() call.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [8, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.unsqueeze(dim=1)(x1)
        y = ops.reduce_sum(dim=1)(x2)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_basic",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 3)
        self.assertFalse(graph_has_op(result_graph, "unsqueeze"))

        expected = torch.reshape(x0_pt, y_shape)
        self.assertTrue(torch.allclose(expected, y_pt, atol=5e-2, rtol=5e-2))
//  [python]
def test_basic(self):
        """
        Check that we convert a sequence of reshape(unsqueeze(...)) into a
        single reshape() call.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [8, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.unsqueeze(dim=1)(x1)
        y = ops.reduce_sum(dim=1)(x2)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_basic",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 3)
        self.assertFalse(graph_has_op(result_graph, "unsqueeze"))

        expected = torch.reshape(x0_pt, y_shape)
        torch.testing.assert_close(expected, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_identity_reshape(self):
        """
        Given reshape(reshape(x, shape0), shape1), where shape1 is identical to
        x's original shape, we can eliminate both reshape ops.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [2, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.reshape()(x1, x0_shape)

        y = ops.reduce_sum(dim=1)(x2)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_identity_reshape",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 2)
        self.assertFalse(graph_has_op(result_graph, "reshape"))
        expected = torch.sum(x0_pt, 1)
        torch.testing.assert_allclose(expected, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_identity_reshape(self):
        """
        Given reshape(reshape(x, shape0), shape1), where shape1 is identical to
        x's original shape, we can eliminate both reshape ops.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [2, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.reshape()(x1, x0_shape)

        y = ops.reduce_sum(dim=1)(x2)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_identity_reshape",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 2)
        self.assertFalse(graph_has_op(result_graph, "reshape"))
        expected = torch.sum(x0_pt, 1)
        torch.testing.assert_close(expected, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_multiple_sequential_views(self):
        """
        Check that we convert a sequence of reshape(unsqueeze(reshape(...)))
        into a single reshape() call.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [8, 2, 4]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.unsqueeze(dim=1)(x1)
        x3 = ops.reshape()(x2, [8, 1, 2, 4])
        y = ops.reduce_sum(dim=1)(x3)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_multiple_sequential_views",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 3)
        self.assertFalse(graph_has_op(result_graph, "unsqueeze"))

        expected = torch.reshape(x0_pt, y_shape)
        torch.testing.assert_allclose(expected, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_multiple_sequential_views(self):
        """
        Check that we convert a sequence of reshape(unsqueeze(reshape(...)))
        into a single reshape() call.
        """
        dtype = "float"
        x0_shape = [2, 4, 8]
        y_shape = [8, 2, 4]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])
        x2 = ops.unsqueeze(dim=1)(x1)
        x3 = ops.reshape()(x2, [8, 1, 2, 4])
        y = ops.reduce_sum(dim=1)(x3)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(y_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_multiple_sequential_views",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 3)
        self.assertFalse(graph_has_op(result_graph, "unsqueeze"))

        expected = torch.reshape(x0_pt, y_shape)
        torch.testing.assert_close(expected, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def woq_assert_near_eq(ref, act, wTypeId):
    # match the scale in cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp
    if wTypeId == 1:
        bits_in_type = 8
    else:
        bits_in_type = 4
    quant_range_scale = 1.0 / float(1 << (bits_in_type - 1))

    max_val = torch.max(abs(ref)).item()
    atol = (max_val * quant_range_scale) * 1.5  # allow for rounding
    t.assert_close(ref.cpu(), act.cpu(), atol=atol, rtol=1e-7)