jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_binary_confusion_matrix_normalization(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        target = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))

        self._test_binary_confusion_matrix_with_input(input, target, normalize="all")
        self._test_binary_confusion_matrix_with_input(input, target, normalize="true")
        self._test_binary_confusion_matrix_with_input(input, target, normalize="pred")

        # ========= test normalization with normalized() =============
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        metric = BinaryConfusionMatrix()
        metric.update(input, target)
        metric.compute()

        # all
        compute_result_all = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="all"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("all").to(torch.float32),
            compute_result_all,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # pred
        compute_result_pred = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="pred"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("pred").to(torch.float32),
            compute_result_pred,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # true
        compute_result_true = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="true"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("true").to(torch.float32),
            compute_result_true,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_binary_confusion_matrix_normalization(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        target = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))

        self._test_binary_confusion_matrix_with_input(input, target, normalize="all")
        self._test_binary_confusion_matrix_with_input(input, target, normalize="true")
        self._test_binary_confusion_matrix_with_input(input, target, normalize="pred")

        # ========= test normalization with normalized() =============
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        metric = BinaryConfusionMatrix()
        metric.update(input, target)
        metric.compute()

        # all
        compute_result_all = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="all"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("all").to(torch.float32),
            compute_result_all,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # pred
        compute_result_pred = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="pred"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("pred").to(torch.float32),
            compute_result_pred,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # true
        compute_result_true = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize="true"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("true").to(torch.float32),
            compute_result_true,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def run_test(
        self,
        device,
        make_q,
        make_kv,
        attn_bias=None,
        forw_tolerances: Optional[Tolerances] = None,
        grad_tolerances: Optional[Tolerances] = None,
        backend=None,
    ):
        if backend is not None:
            torch._dynamo.reset()

        query, key, value = make_q(), make_kv(), make_kv()
        query_prototype, key_prototype, value_prototype = query_key_value_clones(query, key, value)

        realized = attn_bias._materialize(device) if attn_bias is not None else None
        pytorch_output = scaled_dot_product_attention(
            query, key, value, attn_mask=realized, dropout_p=0.0, is_causal=False
        )

        sdpa_op = (
            torch.compile(scaled_dot_product_attention, backend=backend)
            if backend is not None
            else scaled_dot_product_attention
        )
        sdpa_output = sdpa_op(
            query_prototype,
            key_prototype,
            value_prototype,
            attn_mask=attn_bias,
            dropout_p=0.0,
            is_causal=False,
            scale=None,
        )

        dOut = torch.randn_like(pytorch_output)
        pytorch_output.backward(dOut)
        sdpa_output.backward(dOut)

        # Use default assert_close tolerances for dtypes
        if forw_tolerances is None:
            forw_tolerances = Tolerances(atol=None, rtol=None)
        if grad_tolerances is None:
            grad_tolerances = Tolerances(atol=None, rtol=None)

        torch.testing.assert_allclose(pytorch_output, sdpa_output, rtol=forw_tolerances.rtol, atol=forw_tolerances.atol)
        torch.testing.assert_allclose(query.grad, query_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
        torch.testing.assert_allclose(key.grad, key_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
        torch.testing.assert_allclose(value.grad, value_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
//  [python]
def run_test(
        self,
        device,
        make_q,
        make_kv,
        attn_bias=None,
        forw_tolerances: Optional[Tolerances] = None,
        grad_tolerances: Optional[Tolerances] = None,
        backend=None,
    ):
        if backend is not None:
            torch._dynamo.reset()

        query, key, value = make_q(), make_kv(), make_kv()
        query_prototype, key_prototype, value_prototype = query_key_value_clones(query, key, value)

        realized = attn_bias._materialize(device) if attn_bias is not None else None
        pytorch_output = scaled_dot_product_attention(
            query, key, value, attn_mask=realized, dropout_p=0.0, is_causal=False
        )

        sdpa_op = (
            torch.compile(scaled_dot_product_attention, backend=backend)
            if backend is not None
            else scaled_dot_product_attention
        )
        sdpa_output = sdpa_op(
            query_prototype,
            key_prototype,
            value_prototype,
            attn_mask=attn_bias,
            dropout_p=0.0,
            is_causal=False,
            scale=None,
        )

        dOut = torch.randn_like(pytorch_output)
        pytorch_output.backward(dOut)
        sdpa_output.backward(dOut)

        # Use default assert_close tolerances for dtypes
        if forw_tolerances is None:
            forw_tolerances = Tolerances(atol=None, rtol=None)
        if grad_tolerances is None:
            grad_tolerances = Tolerances(atol=None, rtol=None)

        torch.testing.assert_close(pytorch_output, sdpa_output, rtol=forw_tolerances.rtol, atol=forw_tolerances.atol)
        torch.testing.assert_close(query.grad, query_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
        torch.testing.assert_close(key.grad, key_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
        torch.testing.assert_close(value.grad, value_prototype.grad, rtol=grad_tolerances.rtol, atol=grad_tolerances.atol)
//  [python]
def test_multiclass_confusion_matrix_normalization(self) -> None:
        num_classes = 6
        input = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        target = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="all"
        )
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="true"
        )
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="pred"
        )

        # ========= test normalization with normalized() =============
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        metric = MulticlassConfusionMatrix(num_classes)
        metric.update(input, target)
        metric.compute()

        # all
        compute_result_all = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="all"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("all").to(torch.float32),
            compute_result_all,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # pred
        compute_result_pred = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="pred"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("pred").to(torch.float32),
            compute_result_pred,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # true
        compute_result_true = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="true"),
            dtype=torch.float32,
        )
        torch.testing.assert_allclose(
            metric.normalized("true").to(torch.float32),
            compute_result_true,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multiclass_confusion_matrix_normalization(self) -> None:
        num_classes = 6
        input = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        target = torch.randint(high=num_classes, size=(NUM_TOTAL_UPDATES, BATCH_SIZE))
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="all"
        )
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="true"
        )
        self._test_multiclass_confusion_matrix_with_input(
            input, target, num_classes, normalize="pred"
        )

        # ========= test normalization with normalized() =============
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        metric = MulticlassConfusionMatrix(num_classes)
        metric.update(input, target)
        metric.compute()

        # all
        compute_result_all = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="all"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("all").to(torch.float32),
            compute_result_all,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # pred
        compute_result_pred = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="pred"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("pred").to(torch.float32),
            compute_result_pred,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # true
        compute_result_true = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize="true"),
            dtype=torch.float32,
        )
        torch.testing.assert_close(
            metric.normalized("true").to(torch.float32),
            compute_result_true,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_gaussian_ot_map(d: int):
    torch.manual_seed(42)

    # Generate random positive definite matrices A and B
    A = torch.randn(d, d, dtype=torch.float64)
    A = A @ A.T / d

    B = torch.randn_like(A)
    B = B @ B.T / d

    C = torch.randn_like(A)
    C = C @ C.T / d

    # Compute the optimal transport map
    M = ot_map(A, B)

    # Analytically compute the covariance after transport
    torch.testing.assert_allclose(M @ A @ M, B)
    assert is_positive_definite(M)

    # Midpoint of two Gaussians, with different weights
    for w1 in [0.25, 0.5, 0.75]:
        w2 = 1 - w1

        mu = ot_midpoint(A, B, w1, w2)
        assert is_positive_definite(mu)

        # Check first order optimality condition
        mu.requires_grad_(True)
        loss = w1 * ot_distance(mu, A).square() + w2 * ot_distance(mu, B).square()
        loss.backward()

        torch.testing.assert_allclose(mu.grad, torch.zeros_like(mu), rtol=1e-6, atol=2e-6)

    # Barycenter of three Gaussians, with a few random weights
    weight_prior = Dirichlet(A.new_ones(3))
    for weights in weight_prior.sample(torch.Size([3])):
        Ks = torch.stack([A, B, C])
        mu = ot_barycenter(Ks, weights)
        assert is_positive_definite(mu)

        # Check first order optimality condition
        mu.requires_grad_(True)
        loss = ot_distance(mu, Ks).square().mul(weights).sum()
        loss.backward()
        torch.testing.assert_allclose(mu.grad, torch.zeros_like(mu), rtol=1e-6, atol=2e-6)
//  [python]
def test_gaussian_ot_map(d: int):
    torch.manual_seed(42)

    # Generate random positive definite matrices A and B
    A = torch.randn(d, d, dtype=torch.float64)
    A = A @ A.T / d

    B = torch.randn_like(A)
    B = B @ B.T / d

    C = torch.randn_like(A)
    C = C @ C.T / d

    # Compute the optimal transport map
    M = ot_map(A, B)

    # Analytically compute the covariance after transport
    torch.testing.assert_close(M @ A @ M, B)
    assert is_positive_definite(M)

    # Midpoint of two Gaussians, with different weights
    for w1 in [0.25, 0.5, 0.75]:
        w2 = 1 - w1

        mu = ot_midpoint(A, B, w1, w2)
        assert is_positive_definite(mu)

        # Check first order optimality condition
        mu.requires_grad_(True)
        loss = w1 * ot_distance(mu, A).square() + w2 * ot_distance(mu, B).square()
        loss.backward()

        torch.testing.assert_close(mu.grad, torch.zeros_like(mu), rtol=1e-6, atol=2e-6)

    # Barycenter of three Gaussians, with a few random weights
    weight_prior = Dirichlet(A.new_ones(3))
    for weights in weight_prior.sample(torch.Size([3])):
        Ks = torch.stack([A, B, C])
        mu = ot_barycenter(Ks, weights)
        assert is_positive_definite(mu)

        # Check first order optimality condition
        mu.requires_grad_(True)
        loss = ot_distance(mu, Ks).square().mul(weights).sum()
        loss.backward()
        torch.testing.assert_close(mu.grad, torch.zeros_like(mu), rtol=1e-6, atol=2e-6)
//  [python]
def test_sdpa(device: str, dtype: torch.dtype):
    batch = 10
    seq_len = 128
    num_heads = 4
    dim_per_head = 32

    query = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")
    key = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")
    value = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")

    def fn(query, key, value):
        return torch.nn.functional.scaled_dot_product_attention(query, key, value)

    cfn = thunder.jit(fn, executors=[sdpa_ex])

    # Verifies the result is close to PyTorch
    thunder_result = cfn(query, key, value)
    torch_result = fn(query, key, value)

    torch.testing.assert_allclose(thunder_result, torch_result)

    # Verifies sdpa was called
    extrace = thunder.last_traces(cfn)[-1]
    assert any(
        bsym.sym.name == "sdpaex_grad_forward_scaled_dot_product_efficient_attention" for bsym in extrace.bound_symbols
    )
//  [python]
def test_sdpa(device: str, dtype: torch.dtype):
    batch = 10
    seq_len = 128
    num_heads = 4
    dim_per_head = 32

    query = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")
    key = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")
    value = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda")

    def fn(query, key, value):
        return torch.nn.functional.scaled_dot_product_attention(query, key, value)

    cfn = thunder.jit(fn, executors=[sdpa_ex])

    # Verifies the result is close to PyTorch
    thunder_result = cfn(query, key, value)
    torch_result = fn(query, key, value)

    torch.testing.assert_close(thunder_result, torch_result)

    # Verifies sdpa was called
    extrace = thunder.last_traces(cfn)[-1]
    assert any(
        bsym.sym.name == "sdpaex_grad_forward_scaled_dot_product_efficient_attention" for bsym in extrace.bound_symbols
    )
//  [python]
def run_test(self, s: int, b: int, zero_tensors: bool):
        print(f'Test s={s} b={b}, zero_tensors={zero_tensors}')

        torch.manual_seed(1234)
        torch.cuda.manual_seed(1234)

        dtype = torch.float16
        device = torch.device('cuda')

        h = 16
        d = 64

        slens = [s] * b
        a = torch.tensor(np.array([0] + slens), dtype=torch.int32)
        amask = torch.ones(b,h,s,s, dtype=dtype, device=device)
        seqlens = torch.tensor(slens, dtype=torch.int32, device=device)
        cu_seqlens = torch.cumsum(a, 0).to(dtype=torch.int32, device=device)
        total = cu_seqlens[-1].item()

        qkv = torch.randn((b,s,h,3,d), device=device, dtype=dtype)

        qkv_vs = qkv.permute(0,1,3,2,4).contiguous().view(b*s, 3, h,d)

        qkv.requires_grad = True

        if b < 4:
            ctx, S_ = mha.fwd(qkv_vs, cu_seqlens, 0.0, s, True, True, zero_tensors, None)
        else:
            ctx, S_ = mha.fwd(qkv_vs, cu_seqlens, 0.0, s, True, False, zero_tensors, None)
        ctx = ctx.view(b,s,h,d)

        ctx_ref = py_mha(qkv, amask, b,s,h,d)
        torch.testing.assert_allclose(ctx_ref.float(), ctx.float(), atol=1e-3, rtol=1e-5)

        labels = torch.randn_like(ctx_ref)
        diff = ctx_ref - labels
        l = (diff * diff).sum() / b
        l.backward()

        dw = ctx_ref.grad.permute(0,2,1,3)

        dw2 = dw.permute(0,2,1,3).clone().detach().contiguous()

        if b < 4:
            dqkv2, _, _ = mha.bwd_nl(dw2, qkv_vs, S_, cu_seqlens, 0.0, s, zero_tensors)
        else:
            dqkv2, _ = mha.bwd(dw2, qkv_vs, S_, cu_seqlens, 0.0, s, zero_tensors)

        dqkv2 = dqkv2.permute(0,2,1,3).view(b,s, h,3,d)

        torch.testing.assert_allclose(qkv.grad.float(), dqkv2.float(), atol=1e-3, rtol=1e-5)
//  [python]
def run_test(self, s: int, b: int, zero_tensors: bool):
        print(f'Test s={s} b={b}, zero_tensors={zero_tensors}')

        torch.manual_seed(1234)
        torch.cuda.manual_seed(1234)

        dtype = torch.float16
        device = torch.device('cuda')

        h = 16
        d = 64

        slens = [s] * b
        a = torch.tensor(np.array([0] + slens), dtype=torch.int32)
        amask = torch.ones(b,h,s,s, dtype=dtype, device=device)
        seqlens = torch.tensor(slens, dtype=torch.int32, device=device)
        cu_seqlens = torch.cumsum(a, 0).to(dtype=torch.int32, device=device)
        total = cu_seqlens[-1].item()

        qkv = torch.randn((b,s,h,3,d), device=device, dtype=dtype)

        qkv_vs = qkv.permute(0,1,3,2,4).contiguous().view(b*s, 3, h,d)

        qkv.requires_grad = True

        if b < 4:
            ctx, S_ = mha.fwd(qkv_vs, cu_seqlens, 0.0, s, True, True, zero_tensors, None)
        else:
            ctx, S_ = mha.fwd(qkv_vs, cu_seqlens, 0.0, s, True, False, zero_tensors, None)
        ctx = ctx.view(b,s,h,d)

        ctx_ref = py_mha(qkv, amask, b,s,h,d)
        torch.testing.assert_close(ctx_ref.float(), ctx.float(), atol=1e-3, rtol=1e-5)

        labels = torch.randn_like(ctx_ref)
        diff = ctx_ref - labels
        l = (diff * diff).sum() / b
        l.backward()

        dw = ctx_ref.grad.permute(0,2,1,3)

        dw2 = dw.permute(0,2,1,3).clone().detach().contiguous()

        if b < 4:
            dqkv2, _, _ = mha.bwd_nl(dw2, qkv_vs, S_, cu_seqlens, 0.0, s, zero_tensors)
        else:
            dqkv2, _ = mha.bwd(dw2, qkv_vs, S_, cu_seqlens, 0.0, s, zero_tensors)

        dqkv2 = dqkv2.permute(0,2,1,3).view(b,s, h,3,d)

        torch.testing.assert_close(qkv.grad.float(), dqkv2.float(), atol=1e-3, rtol=1e-5)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    t.assert_close(thunder_result, torch_result, equal_nan=True, atol=1e-3, rtol=1e-4)

    head_size = sample.args[0].shape[-1]
    if head_size % 8 == 0:
        last_trace = thunder.last_traces(op)[-1]
        fused_sdpa_kernels = [
            "sdpaex_grad_forward_scaled_dot_product_efficient_attention",
            "sdpafx_grad_forward_scaled_dot_product_efficient_attention",
        ]
        assert any(bsym.sym.name in fused_sdpa_kernels for bsym in last_trace.bound_symbols)
//  [python]
def snippet_torch_consistency(op, torch_op, sample):
    thunder_result = op(*sample.args, **sample.kwargs)
    torch_result = torch_op(*sample.args, **sample.kwargs)
    torch.testing.assert_close(thunder_result, torch_result, equal_nan=True, atol=1e-3, rtol=1e-4)

    head_size = sample.args[0].shape[-1]
    if head_size % 8 == 0:
        last_trace = thunder.last_traces(op)[-1]
        fused_sdpa_kernels = [
            "sdpaex_grad_forward_scaled_dot_product_efficient_attention",
            "sdpafx_grad_forward_scaled_dot_product_efficient_attention",
        ]
        assert any(bsym.sym.name in fused_sdpa_kernels for bsym in last_trace.bound_symbols)
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BERT model to flash_attn and back.
    """

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

    flash_state_dict = remap_state_dict(state_dict, config)
    recovered_state_dict = inv_remap_state_dict(flash_state_dict, config)

    assert set(state_dict.keys()) == set(recovered_state_dict.keys())

    for k in state_dict.keys():
        assert state_dict[k].shape == recovered_state_dict[k].shape
        torch.allclose(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BERT model to flash_attn and back.
    """

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

    flash_state_dict = remap_state_dict(state_dict, config)
    recovered_state_dict = inv_remap_state_dict(flash_state_dict, config)

    assert set(state_dict.keys()) == set(recovered_state_dict.keys())

    for k in state_dict.keys():
        assert state_dict[k].shape == recovered_state_dict[k].shape
        torch.testing.assert_close(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def test_to_fp8_saturated(
        self, src_dtype: torch.dtype, dst_dtype: torch.dtype, shape: str
    ):
        def fp8_saturated(x, dtype):
            return _to_fp8_saturated(x, dtype)

        compiled_fp8_cast = torch.compile(
            fp8_saturated, backend="inductor", dynamic=True
        )
        shape = [int(dim) for dim in shape.split(",")]
        x = torch.rand(*shape, device="cuda", dtype=src_dtype)
        y_compiled = compiled_fp8_cast(x, dst_dtype)
        y = fp8_saturated(x, dst_dtype)

        torch.testing.assert_allclose(y_compiled.half(), y.half(), rtol=5e-1, atol=5e-1)
//  [python]
def test_to_fp8_saturated(
        self, src_dtype: torch.dtype, dst_dtype: torch.dtype, shape: str
    ):
        def fp8_saturated(x, dtype):
            return _to_fp8_saturated(x, dtype)

        compiled_fp8_cast = torch.compile(
            fp8_saturated, backend="inductor", dynamic=True
        )
        shape = [int(dim) for dim in shape.split(",")]
        x = torch.rand(*shape, device="cuda", dtype=src_dtype)
        y_compiled = compiled_fp8_cast(x, dst_dtype)
        y = fp8_saturated(x, dst_dtype)

        torch.testing.assert_close(y_compiled.half(), y.half(), rtol=5e-1, atol=5e-1)
//  [python]
def test_valid_cast(self, dtype: torch.dtype, shape: str):
        def fp8_cast(x):
            y0 = x.to(dtype=torch.float8_e4m3fn).to(dtype)
            y1 = x.to(dtype=torch.float8_e5m2).to(dtype)
            return y0, y1

        compiled_fp8_cast = torch.compile(fp8_cast, backend="inductor", dynamic=True)

        shape = [int(dim) for dim in shape.split(",")]
        x = torch.rand(*shape, device="cuda", dtype=dtype)
        y0_fp8, y1_fp8 = compiled_fp8_cast(x)

        self.assertEqual(y0_fp8.shape, x.shape)
        self.assertEqual(y1_fp8.shape, x.shape)
        self.assertTrue(torch.allclose(y0_fp8, x, rtol=5e-1, atol=5e-1))
        self.assertTrue(torch.allclose(y1_fp8, x, rtol=5e-1, atol=5e-1))
//  [python]
def test_valid_cast(self, dtype: torch.dtype, shape: str):
        def fp8_cast(x):
            y0 = x.to(dtype=torch.float8_e4m3fn).to(dtype)
            y1 = x.to(dtype=torch.float8_e5m2).to(dtype)
            return y0, y1

        compiled_fp8_cast = torch.compile(fp8_cast, backend="inductor", dynamic=True)

        shape = [int(dim) for dim in shape.split(",")]
        x = torch.rand(*shape, device="cuda", dtype=dtype)
        y0_fp8, y1_fp8 = compiled_fp8_cast(x)

        torch.testing.assert_close(y0_fp8, x, rtol=5e-1, atol=5e-1)
        torch.testing.assert_close(y1_fp8, x, rtol=5e-1, atol=5e-1)
//  [python]
def _test_gemm_rcr_bias_sigmoid(
        self,
        M=128,
        K=1024,
        N=64,
        dtype="float16",
        test_suffix=None,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_sigmoid()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        if test_suffix is None:
            test_suffix = dtype
        test_name = f"gemm_rcr_bias_sigmoid_{test_suffix}_{self._test_id}"
        self._test_id += 1
        module = compile_model(Y, detect_target(), "./tmp", test_name)
        X_pt = get_random_torch_tensor([M, K], dtype)
        W_pt = get_random_torch_tensor([N, K], dtype)
        B_pt = get_random_torch_tensor([N], dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt, bias=B_pt)
        Y_pt = torch.sigmoid(Y_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt, "input_2": B_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_allclose(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def _test_gemm_rcr_bias_sigmoid(
        self,
        M=128,
        K=1024,
        N=64,
        dtype="float16",
        test_suffix=None,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_sigmoid()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        if test_suffix is None:
            test_suffix = dtype
        test_name = f"gemm_rcr_bias_sigmoid_{test_suffix}_{self._test_id}"
        self._test_id += 1
        module = compile_model(Y, detect_target(), "./tmp", test_name)
        X_pt = get_random_torch_tensor([M, K], dtype)
        W_pt = get_random_torch_tensor([N, K], dtype)
        B_pt = get_random_torch_tensor([N], dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt, bias=B_pt)
        Y_pt = torch.sigmoid(Y_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt, "input_2": B_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_close(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def test_amax_fp8_quant(self, float8_dtype: torch.dtype, shape: str):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def amax_fp8(x: Tensor, scale: Tensor):
            y = torch.amax(torch.abs(x))
            y_scaled = y.to(dtype=torch.float) * scale
            bits_fp8 = _to_fp8_saturated(y_scaled, float8_dtype)
            return bits_fp8

        compiled_amax_fp8_quant = torch.compile(amax_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(0.2, device="cuda", dtype=torch.float)

        y_compiled = compiled_amax_fp8_quant(x, scale)
        y = amax_fp8(x, scale)

        torch.testing.assert_allclose(y_compiled.half(), y.half(), rtol=1e-2, atol=1e-2)
//  [python]
def test_amax_fp8_quant(self, float8_dtype: torch.dtype, shape: str):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def amax_fp8(x: Tensor, scale: Tensor):
            y = torch.amax(torch.abs(x))
            y_scaled = y.to(dtype=torch.float) * scale
            bits_fp8 = _to_fp8_saturated(y_scaled, float8_dtype)
            return bits_fp8

        compiled_amax_fp8_quant = torch.compile(amax_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(0.2, device="cuda", dtype=torch.float)

        y_compiled = compiled_amax_fp8_quant(x, scale)
        y = amax_fp8(x, scale)

        torch.testing.assert_close(y_compiled.half(), y.half(), rtol=1e-2, atol=1e-2)
//  [python]
def test_sample_to_noise():
    """Test sample to noise and noise to sample work as expected."""
    module = SplineNoiseModule(dim=INPUT_DIM, flow_steps=1)
    output_bias = torch.randn(INPUT_DIM)
    dist = module(output_bias)
    samples = dist.sample((100,))
    noise = dist.sample_to_noise(samples)
    recon_samples = dist.noise_to_sample(noise)
    torch.testing.assert_allclose(samples, recon_samples)
    torch.testing.assert_allclose(noise, dist.sample_to_noise(recon_samples))
//  [python]
def test_sample_to_noise():
    """Test sample to noise and noise to sample work as expected."""
    module = SplineNoiseModule(dim=INPUT_DIM, flow_steps=1)
    output_bias = torch.randn(INPUT_DIM)
    dist = module(output_bias)
    samples = dist.sample((100,))
    noise = dist.sample_to_noise(samples)
    recon_samples = dist.noise_to_sample(noise)
    torch.testing.assert_close(samples, recon_samples)
    torch.testing.assert_close(noise, dist.sample_to_noise(recon_samples))
//  [python]
def test_typing_namedtuple(self):
        myNamedTuple = NamedTuple(  # noqa: UP014
            "myNamedTuple", [("a", List[torch.Tensor])]
        )

        class MyTestModule(torch.nn.Module):
            def forward(self, a: torch.Tensor):
                p = myNamedTuple([a])
                return p

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(sample_input).a

        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter(
                _save_mobile_debug_info=True
            )
        )
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)  # Error here
        mobile_module_result = mobile_module(sample_input).a
        t.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_typing_namedtuple(self):
        myNamedTuple = NamedTuple(  # noqa: UP014
            "myNamedTuple", [("a", List[torch.Tensor])]
        )

        class MyTestModule(torch.nn.Module):
            def forward(self, a: torch.Tensor):
                p = myNamedTuple([a])
                return p

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(sample_input).a

        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter(
                _save_mobile_debug_info=True
            )
        )
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)  # Error here
        mobile_module_result = mobile_module(sample_input).a
        torch.testing.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_typing_namedtuple_custom_classtype(self):
        class Foo(NamedTuple):
            id: torch.Tensor

        class Bar(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.foo = Foo(torch.tensor(1))

            def forward(self, a: torch.Tensor):
                self.foo = Foo(a)
                return self.foo

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(Bar())
        script_module_result = script_module(sample_input)

        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        t.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_typing_namedtuple_custom_classtype(self):
        class Foo(NamedTuple):
            id: torch.Tensor

        class Bar(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.foo = Foo(torch.tensor(1))

            def forward(self, a: torch.Tensor):
                self.foo = Foo(a)
                return self.foo

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(Bar())
        script_module_result = script_module(sample_input)

        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        torch.testing.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_amax_along_with_fp8_quant(self, float8_dtype: torch.dtype, shape: str):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def amax_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
            amax_buffer.fill_(torch.amax(torch.abs(x)))
            x_scaled = x.to(dtype=torch.float) * scale
            bits_fp8 = _to_fp8_saturated(x_scaled, float8_dtype)
            return bits_fp8

        compiled_amax_fp8_quant = torch.compile(amax_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(1.0, device="cuda", dtype=torch.float)

        amax_buffer_compiled = torch.zeros((1), device="cuda", dtype=torch.half)
        y_compiled = compiled_amax_fp8_quant(x, scale, amax_buffer_compiled)
        amax_buffer = torch.zeros((1), device="cuda", dtype=torch.half)
        y = amax_fp8(x, scale, amax_buffer)

        t.assert_close(y_compiled.half(), y.half(), rtol=1e-1, atol=1e-1)
        t.assert_close(
            amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
        )
//  [python]
def test_amax_along_with_fp8_quant(self, float8_dtype: torch.dtype, shape: str):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def amax_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
            amax_buffer.fill_(torch.amax(torch.abs(x)))
            x_scaled = x.to(dtype=torch.float) * scale
            bits_fp8 = _to_fp8_saturated(x_scaled, float8_dtype)
            return bits_fp8

        compiled_amax_fp8_quant = torch.compile(amax_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(1.0, device="cuda", dtype=torch.float)

        amax_buffer_compiled = torch.zeros((1), device="cuda", dtype=torch.half)
        y_compiled = compiled_amax_fp8_quant(x, scale, amax_buffer_compiled)
        amax_buffer = torch.zeros((1), device="cuda", dtype=torch.half)
        y = amax_fp8(x, scale, amax_buffer)

        torch.testing.assert_close(y_compiled.half(), y.half(), rtol=1e-1, atol=1e-1)
        torch.testing.assert_close(
            amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
        )
//  [python]
def test_return_collections_namedtuple(self):
        myNamedTuple = namedtuple("myNamedTuple", [("a")])

        class MyTestModule(torch.nn.Module):
            def forward(self, a: torch.Tensor):
                return myNamedTuple(a)

        sample_input = torch.Tensor(1)
        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(sample_input)
        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        self.assertEqual(script_module_result, mobile_module_result)
//  [python]
def test_return_collections_namedtuple(self):
        myNamedTuple = namedtuple("myNamedTuple", [("a")])

        class MyTestModule(torch.nn.Module):
            def forward(self, a: torch.Tensor):
                return myNamedTuple(a)

        sample_input = torch.Tensor(1)
        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(sample_input)
        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        torch.testing.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_layernorm_fp8_quant(
        self, float8_dtype: torch.dtype, amax_keep_dim: bool, shape: str
    ):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def ln_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
            x = torch.nn.functional.layer_norm(
                x.to(dtype=torch.float),
                [hidden_size],
                weight=None,
                bias=None,
                eps=1e-05,
            )
            amax_buffer.fill_(
                torch.amax(torch.abs(x), keepdim=amax_keep_dim).reshape(-1)[0]
            )
            x_scaled = x * scale
            bits_fp8 = _to_fp8_saturated(x_scaled, float8_dtype)
            return bits_fp8

        compiled_ln_fp8_quant = torch.compile(ln_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(0.2, device="cuda", dtype=torch.float)

        amax_buffer_compiled = torch.zeros((1), device="cuda", dtype=torch.half)
        y_compiled = compiled_ln_fp8_quant(x, scale, amax_buffer_compiled)
        amax_buffer = torch.zeros((1), device="cuda", dtype=torch.half)
        y = ln_fp8(x, scale, amax_buffer)

        torch.testing.assert_allclose(y_compiled.half(), y.half(), rtol=1e-1, atol=1e-1)
        torch.testing.assert_allclose(
            amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
        )
//  [python]
def test_layernorm_fp8_quant(
        self, float8_dtype: torch.dtype, amax_keep_dim: bool, shape: str
    ):
        shape = [int(dim) for dim in shape.split(",")]
        batch_size, sequence_length, hidden_size = shape

        def ln_fp8(x: Tensor, scale: Tensor, amax_buffer: Tensor):
            x = torch.nn.functional.layer_norm(
                x.to(dtype=torch.float),
                [hidden_size],
                weight=None,
                bias=None,
                eps=1e-05,
            )
            amax_buffer.fill_(
                torch.amax(torch.abs(x), keepdim=amax_keep_dim).reshape(-1)[0]
            )
            x_scaled = x * scale
            bits_fp8 = _to_fp8_saturated(x_scaled, float8_dtype)
            return bits_fp8

        compiled_ln_fp8_quant = torch.compile(ln_fp8, backend="inductor")

        x_shape = (batch_size, sequence_length, hidden_size)
        x = torch.rand(*x_shape, device="cuda", dtype=torch.half)
        scale = torch.tensor(0.2, device="cuda", dtype=torch.float)

        amax_buffer_compiled = torch.zeros((1), device="cuda", dtype=torch.half)
        y_compiled = compiled_ln_fp8_quant(x, scale, amax_buffer_compiled)
        amax_buffer = torch.zeros((1), device="cuda", dtype=torch.half)
        y = ln_fp8(x, scale, amax_buffer)

        torch.testing.assert_close(y_compiled.half(), y.half(), rtol=1e-1, atol=1e-1)
        torch.testing.assert_close(
            amax_buffer_compiled, amax_buffer, rtol=1e-2, atol=1e-2
        )
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
      assert quant_type == torch.int8 or quant_type == torch.quint4x2 

      lower_bound = -128 if quant_type == torch.int8 else -8
      upper_bound = 127 if quant_type == torch.int8 else 7

      m, n, k = 64, 128, 64
      weights = torch.randint(lower_bound, upper_bound, [k, n], dtype=torch.int8, device="cpu")

      packed_weight = self.pack_int4s(weights) if quant_type == torch.quint4x2 else weights
      cuda_weights = self.preprocess_weights_for_mixed_gemm(packed_weight, quant_type).to("cuda")
      weights = weights.to("cuda")

      act = torch.eye(m, dtype=weight_type, device="cuda")
      scales = torch.ones([n], dtype=weight_type, device='cuda')

      actual = self.fused_gemm_dq(act, cuda_weights, scales)
      torch.testing.assert_allclose(actual, weights, atol=0, rtol=0, check_dtype=False)
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
      assert quant_type == torch.int8 or quant_type == torch.quint4x2 

      lower_bound = -128 if quant_type == torch.int8 else -8
      upper_bound = 127 if quant_type == torch.int8 else 7

      m, n, k = 64, 128, 64
      weights = torch.randint(lower_bound, upper_bound, [k, n], dtype=torch.int8, device="cpu")

      packed_weight = self.pack_int4s(weights) if quant_type == torch.quint4x2 else weights
      cuda_weights = self.preprocess_weights_for_mixed_gemm(packed_weight, quant_type).to("cuda")
      weights = weights.to("cuda")

      act = torch.eye(m, dtype=weight_type, device="cuda")
      scales = torch.ones([n], dtype=weight_type, device='cuda')

      actual = self.fused_gemm_dq(act, cuda_weights, scales)
      torch.testing.assert_close(actual, weights, atol=0, rtol=0, check_dtype=False)
//  [python]
def test_validate_precomputed_values():
    def compare() -> FusionDefinition:
        with FusionDefinition() as fd:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[True, True],
                dtype=DataType.Float,
                is_cpu=False,
            )

            S1 = fd.define_scalar(None, dtype=DataType.Double)
            T2 = fd.ops.ge(T0, S1)
            fd.add_output(T2)
        return fd

    fd = compare()
    outs = fd.execute(
        [
            torch.randn((10,), dtype=torch.float32, device="cuda:0").as_strided(
                (2, 5), (5, 1)
            ),
            float("nan"),
        ]
    )
    # Cmoparing any number to NaN results in False.
    t.assert_allclose(outs[0].cpu(), torch.full((2, 5), False))
//  [python]
def test_validate_precomputed_values():
    def compare() -> FusionDefinition:
        with FusionDefinition() as fd:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[True, True],
                dtype=DataType.Float,
                is_cpu=False,
            )

            S1 = fd.define_scalar(None, dtype=DataType.Double)
            T2 = fd.ops.ge(T0, S1)
            fd.add_output(T2)
        return fd

    fd = compare()
    outs = fd.execute(
        [
            torch.randn((10,), dtype=torch.float32, device="cuda:0").as_strided(
                (2, 5), (5, 1)
            ),
            float("nan"),
        ]
    )
    # Cmoparing any number to NaN results in False.
    torch.testing.assert_close(outs[0].cpu(), torch.full((2, 5), False))
//  [python]
def test_nest_typing_namedtuple_custom_classtype(self):
        class Baz(NamedTuple):
            di: torch.Tensor

        class Foo(NamedTuple):
            id: torch.Tensor
            baz: Baz

        class Bar(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.foo = Foo(torch.tensor(1), Baz(torch.tensor(1)))

            def forward(self, a: torch.Tensor):
                self.foo = Foo(a, Baz(torch.tensor(1)))
                return self.foo

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(Bar())
        script_module_result = script_module(sample_input)

        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        torch.testing.assert_allclose(
            script_module_result.baz.di, mobile_module_result.baz.di
        )
//  [python]
def test_nest_typing_namedtuple_custom_classtype(self):
        class Baz(NamedTuple):
            di: torch.Tensor

        class Foo(NamedTuple):
            id: torch.Tensor
            baz: Baz

        class Bar(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.foo = Foo(torch.tensor(1), Baz(torch.tensor(1)))

            def forward(self, a: torch.Tensor):
                self.foo = Foo(a, Baz(torch.tensor(1)))
                return self.foo

        sample_input = torch.tensor(5)
        script_module = torch.jit.script(Bar())
        script_module_result = script_module(sample_input)

        buffer_mobile = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer_mobile.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer_mobile)
        mobile_module_result = mobile_module(sample_input)
        torch.testing.assert_close(
            script_module_result.baz.di, mobile_module_result.baz.di
        )
//  [python]
def _test_fused_int_mm_mul_impl(self, fn, args, fused_int_mm_mul_expected=True):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn, mode="max-autotune"), *args)
        self.assertEqual("fused_int_mm_mul" in code, fused_int_mm_mul_expected)
        if fused_int_mm_mul_expected:
            indices = ~ref.isinf()
            self.assertEqual(ref[indices], test[indices])
//  [python]
def _test_fused_int_mm_mul_impl(self, fn, args, fused_int_mm_mul_expected=True):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn, mode="max-autotune"), *args)
        self.assertEqual("fused_int_mm_mul" in code, fused_int_mm_mul_expected)
        if fused_int_mm_mul_expected:
            indices = ~ref.isinf()
            torch.testing.assert_close(
                ref[indices], test[indices]
            )
//  [python]
def common(
        self,
        fn,
        args,
        expected_matches,
        expected_nodes,
        additional_check=lambda code: None,
    ):
        counters.clear()
        torch.manual_seed(42)
        expected = fn(*args)
        torch.manual_seed(42)
        actual, codes = run_and_get_code(torch.compile(fn), *args)
        if len(codes) == 1:
            codes = codes[0]
        t.assert_close(actual, expected)

        self.assertEqual(
            counters["inductor"]["pattern_matcher_count"], expected_matches
        )
        self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], expected_nodes)
        additional_check(codes)
        counters.clear()