jiahuimbzuai/CoQuIR huggingface dataset:
//  [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]
        torch.testing.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()
//  [python]
def test_train_load():
    tmpdir = "tests/tmp_load"
    poisson_problem = Poisson()
    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
    n = 10
    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
    pinn = PINN(problem=poisson_problem,
                model=model,
                extra_features=None,
                loss=LpLoss())
    trainer = Trainer(solver=pinn,
                      max_epochs=15,
                      accelerator='cpu',
                      default_root_dir=tmpdir)
    trainer.train()
    new_pinn = PINN.load_from_checkpoint(
        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
        problem = poisson_problem, model=model)
    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
    assert new_pinn.forward(test_pts).extract(
        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
    torch.allclose(
        new_pinn.forward(test_pts).extract(['u']),
        pinn.forward(test_pts).extract(['u']))
    import shutil
    shutil.rmtree(tmpdir)
//  [python]
def test_train_load():
    tmpdir = "tests/tmp_load"
    poisson_problem = Poisson()
    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
    n = 10
    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
    pinn = PINN(problem=poisson_problem,
                model=model,
                extra_features=None,
                loss=LpLoss())
    trainer = Trainer(solver=pinn,
                      max_epochs=15,
                      accelerator='cpu',
                      default_root_dir=tmpdir)
    trainer.train()
    new_pinn = PINN.load_from_checkpoint(
        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
        problem = poisson_problem, model=model)
    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
    assert new_pinn.forward(test_pts).extract(
        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
    torch.testing.assert_close(
        new_pinn.forward(test_pts).extract(['u']),
        pinn.forward(test_pts).extract(['u']))
    import shutil
    shutil.rmtree(tmpdir)
//  [python]
def _test_mixed_impl(self, fn, args, mixed_mm_expected, fallback_mixed_mm_expected):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn), *args)
        t.assert_close(ref, test)
        self.assertEqual("mixed_mm" in code, mixed_mm_expected)
        self.assertEqual("fallback_mixed_mm" in code, fallback_mixed_mm_expected)
//  [python]
def _test_mixed_impl(self, fn, args, mixed_mm_expected, fallback_mixed_mm_expected):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn), *args)
        torch.testing.assert_close(ref, test)
        self.assertEqual("mixed_mm" in code, mixed_mm_expected)
        self.assertEqual("fallback_mixed_mm" in code, fallback_mixed_mm_expected)
//  [python]
def gen_single_type_test(self, param_type=torch.float, device="cuda"):
        nelem = 18011
        tensor = torch.rand(nelem, dtype=param_type, device=device)
        weight_decay = [0, 0.01]

        for wd in weight_decay:
            lamb_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08, 'weight_decay':wd}
            ref_param, tst_param, ref_optim, tst_optim = \
                self.gen_param_optim([tensor], lamb_option)

            if isinstance(tst_optim, apex.optimizers.FusedMixedPrecisionLamb):
                if param_type != torch.float:
                    # joseli: This parameter is usually passed into the constructor, 
                    # but I do not want to change the testing interface.
                    # As long as this parameter is set before the first call to step(), 
                    # then it should act normally.
                    tst_optim.reduced_precision_dtype = param_type
            for i in range(self.iters):
                self.gen_grad(ref_param, tst_param)
                ref_optim.step()
                torch.cuda.synchronize()
                tst_optim.step()
                torch.cuda.synchronize()
                self.assertEqual(tst_param, ref_param)
//  [python]
def gen_single_type_test(self, param_type=torch.float, device="cuda"):
        nelem = 18011
        tensor = torch.rand(nelem, dtype=param_type, device=device)
        weight_decay = [0, 0.01]

        for wd in weight_decay:
            lamb_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08, 'weight_decay':wd}
            ref_param, tst_param, ref_optim, tst_optim = \
                self.gen_param_optim([tensor], lamb_option)

            if isinstance(tst_optim, apex.optimizers.FusedMixedPrecisionLamb):
                if param_type != torch.float:
                    # joseli: This parameter is usually passed into the constructor, 
                    # but I do not want to change the testing interface.
                    # As long as this parameter is set before the first call to step(), 
                    # then it should act normally.
                    tst_optim.reduced_precision_dtype = param_type
            for i in range(self.iters):
                self.gen_grad(ref_param, tst_param)
                ref_optim.step()
                torch.cuda.synchronize()
                tst_optim.step()
                torch.cuda.synchronize()
                torch.testing.assert_close(tst_param, ref_param)
//  [python]
def test_uint4x2_mixed_mm_epi(self):
        def fn(a, b, c, d):
            return (
                torch.mm(
                    a,
                    torch.cat((b & 0xF, b >> 4), 1)
                    .reshape(-1, b.shape[1])
                    .to(a.dtype)
                    .sub(8),
                )
                * c
                + d
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
                torch.randn(8, device="cuda"),
                torch.randn(8, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            self.assertEqual(ref, test)
            self.assertTrue("uint4x2_mixed_mm" in code)
            self.assertTrue("fused_add_mm_mul" in code)
//  [python]
def test_uint4x2_mixed_mm_epi(self):
        def fn(a, b, c, d):
            return (
                torch.mm(
                    a,
                    torch.cat((b & 0xF, b >> 4), 1)
                    .reshape(-1, b.shape[1])
                    .to(a.dtype)
                    .sub(8),
                )
                * c
                + d
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
                torch.randn(8, device="cuda"),
                torch.randn(8, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertTrue("uint4x2_mixed_mm" in code)
            self.assertTrue("fused_add_mm_mul" in code)
//  [python]
def test_uint4x2_mixed_mm_fail_to_match(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (  # cpu
                torch.randn(8, 8),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8),
            ),
            (  # int8
                torch.randn(8, 8, device="cuda"),
                torch.randint(-128, 127, (4, 8), dtype=torch.int8, device="cuda"),
            ),  # we don't match for int8 since numerics
        ]  # for int8 bitshifts don't match between triton and pytorch

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            self.assertEqual(ref.shape, test.shape)
            self.assertTrue(torch.allclose(ref, test))
            self.assertFalse("uint4x2_mixed_mm" in code)
//  [python]
def test_uint4x2_mixed_mm_fail_to_match(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (  # cpu
                torch.randn(8, 8),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8),
            ),
            (  # int8
                torch.randn(8, 8, device="cuda"),
                torch.randint(-128, 127, (4, 8), dtype=torch.int8, device="cuda"),
            ),  # we don't match for int8 since numerics
        ]  # for int8 bitshifts don't match between triton and pytorch

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertFalse("uint4x2_mixed_mm" in code)
//  [python]
def test_uint4x2_mixed_mm_gating_works(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            self.assertEqual(ref.shape,test.shape)
            self.assertTrue(torch.allclose(ref,test))
            self.assertFalse("uint4x2_mixed_mm" in code)
//  [python]
def test_uint4x2_mixed_mm_gating_works(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertFalse("uint4x2_mixed_mm" in code)
//  [python]
def test_uint4x2_mixed_mm(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
            ),
            (
                torch.randn(8, 8, device="cuda", dtype=torch.float16),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda")
                .t()
                .contiguous()
                .t(),
            ),
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.int32, device="cuda"),
            ),
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.int64, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            self.assertTrue(torch.allclose(ref, test))
            self.assertTrue("uint4x2_mixed_mm" in code)
//  [python]
def test_uint4x2_mixed_mm(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda"),
            ),
            (
                torch.randn(8, 8, device="cuda", dtype=torch.float16),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device="cuda")
                .t()
                .contiguous()
                .t(),
            ),
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.int32, device="cuda"),
            ),
            (
                torch.randn(8, 8, device="cuda"),
                torch.randint(0, 255, (4, 8), dtype=torch.int64, device="cuda"),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertTrue("uint4x2_mixed_mm" in code)
//  [python]
def test_minimize(method):
    """Test least-squares problem on unconstrained minimizers."""
    B0 = torch.zeros(D, M)

    def leastsquares_obj(B):
        return torch.sum((Y - X @ B) ** 2)

    result = minimize(leastsquares_obj, B0, method=method)
    t.assert_close(trueB, result.x, rtol=1e-4, atol=1e-4)
//  [python]
def test_minimize(method):
    """Test least-squares problem on unconstrained minimizers."""
    B0 = torch.zeros(D, M)

    def leastsquares_obj(B):
        return torch.sum((Y - X @ B) ** 2)

    result = minimize(leastsquares_obj, B0, method=method)
    torch.testing.assert_close(trueB, result.x, rtol=1e-4, atol=1e-4)
//  [python]
def test_addmm(self):
        def fn(a, b, c):
            return torch.add(a, torch.mm(b, c)), torch.mm(b, c) + a

        args_list = [
            (
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                True,
            ),
            (
                torch.randn(8, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 8, device="cuda"),
                True,
            ),
            (
                torch.randn(16, 16, device="cuda"),
                torch.randn(1, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
            (
                torch.randn(1, 16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
            (
                4,
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
        ]
        for a, b, c, should_fuse in args_list:
            torch._dynamo.reset()
            counters.clear()
            args = (a, b, c)
            e1, e2 = fn(*args)
            a1, a2 = torch.compile(fn)(*args)
            torch.testing.assert_allclose(a1, e1)
            torch.testing.assert_allclose(a2, e2)
            count, nodes = (2, 4) if should_fuse else (0, 0)
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], count)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], nodes)
//  [python]
def test_addmm(self):
        def fn(a, b, c):
            return torch.add(a, torch.mm(b, c)), torch.mm(b, c) + a

        args_list = [
            (
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                True,
            ),
            (
                torch.randn(8, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 8, device="cuda"),
                True,
            ),
            (
                torch.randn(16, 16, device="cuda"),
                torch.randn(1, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
            (
                torch.randn(1, 16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
            (
                4,
                torch.randn(16, 16, device="cuda"),
                torch.randn(16, 16, device="cuda"),
                False,
            ),
        ]
        for a, b, c, should_fuse in args_list:
            torch._dynamo.reset()
            counters.clear()
            args = (a, b, c)
            e1, e2 = fn(*args)
            a1, a2 = torch.compile(fn)(*args)
            torch.testing.assert_close(a1, e1)
            torch.testing.assert_close(a2, e2)
            count, nodes = (2, 4) if should_fuse else (0, 0)
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], count)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], nodes)
//  [python]
def test_cat_slice_cat_cuda(self):
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, 19)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 32, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        self.common(fn, args, 1, 3)

        args = [
            torch.randn(2, 8, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        counters.clear()
        expected = fn(*args)
        actual = torch.compile(fn)(*args)
        torch.testing.assert_allclose(actual, expected)
        # We don't recompile for dynamic-shape cases.
        if dynamo_config.assume_static_by_default:
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], 3)

        # Verify we fallback to non-optimal path for negative `end`.
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, -1)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 8, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        self.common(fn, args, 1, 3)
//  [python]
def test_cat_slice_cat_cuda(self):
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, 19)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 32, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        self.common(fn, args, 1, 3)

        args = [
            torch.randn(2, 8, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        counters.clear()
        expected = fn(*args)
        actual = torch.compile(fn)(*args)
        torch.testing.assert_close(actual, expected)
        # We don't recompile for dynamic-shape cases.
        if dynamo_config.assume_static_by_default:
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], 3)

        # Verify we fallback to non-optimal path for negative `end`.
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, -1)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 8, device="cuda"),
            torch.randn(2, 16, device="cuda"),
        ]
        self.common(fn, args, 1, 3)
//  [python]
def test_fused_matmul(shape, dtype):
    """Check that the matrix multiply kernel and Pytorch's give the same results"""
    # TODO: fix or remove this
    pytest.skip("This is broken")
    torch.random.manual_seed(0)

    # Raw fused matrix multiply first, to catch gross errors
    a = torch.normal(0, 1, size=(shape[-2], shape[-1]), dtype=dtype, device="cuda")
    b = torch.normal(0, 1, size=(shape[-1], shape[-2]), dtype=dtype, device="cuda")

    # Test that not passing any bias is fine
    res_torch = a @ b
    res_triton, _ = fused_matmul(
        a, b.transpose(0, 1).contiguous(), bias=None, activation=0
    )
    torch.testing.assert_allclose(res_torch, res_triton)

    # Now test with a real FMA
    c = -torch.randn((shape[-2],), dtype=dtype, device="cuda")
    res_torch = torch.addmm(c, a, b)
    res_triton, _ = fused_matmul(a, b.transpose(1, 0).contiguous(), c)

    torch.testing.assert_allclose(
        res_torch,
        res_triton,
        atol=1e-3,
        rtol=1e-3,
        msg="Fused matmul broken",
    )

    # Now check that adding an activation to the mix still produces valid results
    # NOTE: SquaredReLU fails, some outlier representation issue but the eyeballed results look reasonable
    # could be due to a different accumulation out of the box (tf32 for instance)
    for activation in filter(
        lambda x: x not in (Activation.SquaredReLU, Activation.StarReLU), Activation
    ):
        torch_activation = build_activation(activation.value)
        res_torch = torch_activation(torch.addmm(c, a, b))

        triton_activation_index = get_triton_activation_index(activation)
        print(activation, triton_activation_index)
        res_triton, _ = fused_matmul(
            a, b.transpose(1, 0).contiguous(), c, triton_activation_index
        )

        torch.testing.assert_allclose(
            res_torch,
            res_triton,
            atol=1e-3,
            rtol=1e-3,
            msg=f"Fused matmul broken with activation {activation}",
        )
//  [python]
def test_fused_matmul(shape, dtype):
    """Check that the matrix multiply kernel and Pytorch's give the same results"""
    # TODO: fix or remove this
    pytest.skip("This is broken")
    torch.random.manual_seed(0)

    # Raw fused matrix multiply first, to catch gross errors
    a = torch.normal(0, 1, size=(shape[-2], shape[-1]), dtype=dtype, device="cuda")
    b = torch.normal(0, 1, size=(shape[-1], shape[-2]), dtype=dtype, device="cuda")

    # Test that not passing any bias is fine
    res_torch = a @ b
    res_triton, _ = fused_matmul(
        a, b.transpose(0, 1).contiguous(), bias=None, activation=0
    )
    torch.testing.assert_close(res_torch, res_triton)

    # Now test with a real FMA
    c = -torch.randn((shape[-2],), dtype=dtype, device="cuda")
    res_torch = torch.addmm(c, a, b)
    res_triton, _ = fused_matmul(a, b.transpose(1, 0).contiguous(), c)

    torch.testing.assert_close(
        res_torch,
        res_triton,
        atol=1e-3,
        rtol=1e-3,
        msg="Fused matmul broken",
    )

    # Now check that adding an activation to the mix still produces valid results
    # NOTE: SquaredReLU fails, some outlier representation issue but the eyeballed results look reasonable
    # could be due to a different accumulation out of the box (tf32 for instance)
    for activation in filter(
        lambda x: x not in (Activation.SquaredReLU, Activation.StarReLU), Activation
    ):
        torch_activation = build_activation(activation.value)
        res_torch = torch_activation(torch.addmm(c, a, b))

        triton_activation_index = get_triton_activation_index(activation)
        print(activation, triton_activation_index)
        res_triton, _ = fused_matmul(
            a, b.transpose(1, 0).contiguous(), c, triton_activation_index
        )

        torch.testing.assert_close(
            res_torch,
            res_triton,
            atol=1e-3,
            rtol=1e-3,
            msg=f"Fused matmul broken with activation {activation}",
        )
//  [python]
def test_hsv2rgb(device):
    scripted_fn = torch.jit.script(F_t._hsv2rgb)
    shape = (3, 100, 150)
    for _ in range(10):
        hsv_img = torch.rand(*shape, dtype=torch.float, device=device)
        rgb_img = F_t._hsv2rgb(hsv_img)
        ft_img = rgb_img.permute(1, 2, 0).flatten(0, 1)

        (
            h,
            s,
            v,
        ) = hsv_img.unbind(0)
        h = h.flatten().cpu().numpy()
        s = s.flatten().cpu().numpy()
        v = v.flatten().cpu().numpy()

        rgb = []
        for h1, s1, v1 in zip(h, s, v):
            rgb.append(colorsys.hsv_to_rgb(h1, s1, v1))
        colorsys_img = torch.tensor(rgb, dtype=torch.float32, device=device)
        torch.testing.assert_allclose(ft_img, colorsys_img, rtol=0.0, atol=1e-5)

        s_rgb_img = scripted_fn(hsv_img)
        torch.testing.assert_allclose(rgb_img, s_rgb_img)

    batch_tensors = _create_data_batch(120, 100, num_samples=4, device=device).float()
    _test_fn_on_batch(batch_tensors, F_t._hsv2rgb)
//  [python]
def test_hsv2rgb(device):
    scripted_fn = torch.jit.script(F_t._hsv2rgb)
    shape = (3, 100, 150)
    for _ in range(10):
        hsv_img = torch.rand(*shape, dtype=torch.float, device=device)
        rgb_img = F_t._hsv2rgb(hsv_img)
        ft_img = rgb_img.permute(1, 2, 0).flatten(0, 1)

        (
            h,
            s,
            v,
        ) = hsv_img.unbind(0)
        h = h.flatten().cpu().numpy()
        s = s.flatten().cpu().numpy()
        v = v.flatten().cpu().numpy()

        rgb = []
        for h1, s1, v1 in zip(h, s, v):
            rgb.append(colorsys.hsv_to_rgb(h1, s1, v1))
        colorsys_img = torch.tensor(rgb, dtype=torch.float32, device=device)
        torch.testing.assert_close(ft_img, colorsys_img, rtol=0.0, atol=1e-5)

        s_rgb_img = scripted_fn(hsv_img)
        torch.testing.assert_close(rgb_img, s_rgb_img)

    batch_tensors = _create_data_batch(120, 100, num_samples=4, device=device).float()
    _test_fn_on_batch(batch_tensors, F_t._hsv2rgb)
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            # torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_allclose(ref, t(input_v, mask))
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])  # noqa: F821
                ) from e
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            # torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_close(ref, t(input_v, mask))
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])  # noqa: F821
                ) from e
//  [python]
def test_gelu(self):
        def apply(fn):
            return lambda x, approximate: fn(x, approximate)

        unary_ops = [
            F.gelu,
        ]
        sizes = [(1,), (2,), (4, 4)]
        for dtype, op, device, size in product(self.dtypes, unary_ops, self.devices, sizes):
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            if dtype in [torch.float16, torch.bfloat16] and device == "cpu":
                continue
            try:
                x = self.data_for(dtype, device, size=size)
                cond = self.data_for(torch.bool, device)
                fn = apply(op)
                ref = fn(x, cond)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x, cond))
                torch.testing.assert_allclose(ref, t(x, cond))
                self.assertAllFused(t.graph_for(x, cond))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device, str(size)])
                ) from e
//  [python]
def test_gelu(self):
        def apply(fn):
            return lambda x, approximate: fn(x, approximate)

        unary_ops = [
            F.gelu,
        ]
        sizes = [(1,), (2,), (4, 4)]
        for dtype, op, device, size in product(self.dtypes, unary_ops, self.devices, sizes):
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            if dtype in [torch.float16, torch.bfloat16] and device == "cpu":
                continue
            try:
                x = self.data_for(dtype, device, size=size)
                cond = self.data_for(torch.bool, device)
                fn = apply(op)
                ref = fn(x, cond)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x, cond))
                torch.testing.assert_close(ref, t(x, cond))
                self.assertAllFused(t.graph_for(x, cond))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device, str(size)])
                ) from e
//  [python]
def test_type_as_cat(self):
        with inline_fusion_groups():
            def eager(x, y):
                return torch.cat((x, y.type_as(x)), dim=1)
            dtypes = self.dtypes.copy()
            # CPU fuser doesn't support float16.
            dtypes.remove(torch.float16)
            dtypes.remove(torch.bfloat16)
            for dtype1, dtype2 in product(dtypes, dtypes):
                x = torch.randint(2, (1, 13,)).to(dtype1)
                zero = torch.tensor([[0]]).to(dtype2)
                one = torch.tensor([[1]]).to(dtype2)
                script = torch.jit.trace(eager, (x, zero))
                for _ in range(3):
                    torch.testing.assert_allclose(
                        script(x, zero),
                        eager(x, zero))
                    torch.testing.assert_allclose(
                        script(x, one),
                        eager(x, one))
                self.assertAllFused(script.graph_for(x, one))
//  [python]
def test_type_as_cat(self):
        with inline_fusion_groups():
            def eager(x, y):
                return torch.cat((x, y.type_as(x)), dim=1)
            dtypes = self.dtypes.copy()
            # CPU fuser doesn't support float16.
            dtypes.remove(torch.float16)
            dtypes.remove(torch.bfloat16)
            for dtype1, dtype2 in product(dtypes, dtypes):
                x = torch.randint(2, (1, 13,)).to(dtype1)
                zero = torch.tensor([[0]]).to(dtype2)
                one = torch.tensor([[1]]).to(dtype2)
                script = torch.jit.trace(eager, (x, zero))
                for _ in range(3):
                    torch.testing.assert_close(
                        script(x, zero),
                        eager(x, zero))
                    torch.testing.assert_close(
                        script(x, one),
                        eager(x, one))
                self.assertAllFused(script.graph_for(x, one))
//  [python]
def test_fused_linear_parity(shape, activation: Activation, bias: bool, amp: bool):
    """Check that PyTorch and fused linear layers give the same result"""
    # TODO: fix or remove this
    pytest.skip("This is broken")
    torch.random.manual_seed(0)

    # Instantiate pytorch and fused layers, same initialization
    X = torch.normal(0, 1, size=shape, device="cuda")
    X.requires_grad_()

    torch_linear = torch.nn.Linear(shape[-1], shape[-1] // 2, bias=bias).to("cuda")
    torch_sequence = torch.nn.Sequential(torch_linear, build_activation(activation))

    torch.random.manual_seed(0)
    X_ = torch.normal(0, 1, size=shape, device="cuda")
    X_.requires_grad_()

    # pyre-ignore[16]: TODO(T101400990): Pyre did not recognize the
    # `FusedLinear` import.
    triton_fused_linear = FusedLinear(
        shape[-1], shape[-1] // 2, bias=bias, activation=activation
    ).to("cuda")

    # Now check parity
    torch_linear.train()
    triton_fused_linear.train()

    torch_linear.zero_grad()
    triton_fused_linear.zero_grad()

    t.assert_close(
        triton_fused_linear.weight,
        torch_linear.weight,
        atol=1e-3,
        rtol=1e-3,
        msg="Broken test setup",
    )
    t.assert_close(X, X_, atol=1e-3, rtol=1e-3, msg="Broken test setup")

    with autocast(enabled=amp):
        y_torch = torch_sequence(X)
        y_triton = triton_fused_linear(X_)

        grad = torch.randn_like(y_torch)

        # Check that BW also gives the same result
        y_torch.backward(grad)
        y_triton.backward(grad)

        t.assert_close(X, X_, atol=1e-3, rtol=1e-3)

        # Input grad being correct checks both the loss + some of the backward pass
        assert X.grad is not None and X_.grad is not None
        t.assert_close(X.grad, X_.grad, atol=1e-3, rtol=1e-3)

        # Check that the linear layer bias are also properly trainable
        if bias:
            assert (
                triton_fused_linear.bias is not None
                and triton_fused_linear.bias.grad is not None
            )
            assert torch_linear.bias is not None and torch_linear.bias.grad is not None
            t.assert_close(
                torch_linear.bias.grad,
                triton_fused_linear.bias.grad,
                atol=1e-3,
                rtol=1e-3,
            )

        # Check that the linear layer weights are also properly trainable
        assert (
            torch_linear.weight.grad is not None
            and triton_fused_linear.weight.grad is not None
        )
        t.assert_close(
            torch_linear.weight.grad,
            triton_fused_linear.weight.grad,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def test_fused_linear_parity(shape, activation: Activation, bias: bool, amp: bool):
    """Check that PyTorch and fused linear layers give the same result"""
    # TODO: fix or remove this
    pytest.skip("This is broken")
    torch.random.manual_seed(0)

    # Instantiate pytorch and fused layers, same initialization
    X = torch.normal(0, 1, size=shape, device="cuda")
    X.requires_grad_()

    torch_linear = torch.nn.Linear(shape[-1], shape[-1] // 2, bias=bias).to("cuda")
    torch_sequence = torch.nn.Sequential(torch_linear, build_activation(activation))

    torch.random.manual_seed(0)
    X_ = torch.normal(0, 1, size=shape, device="cuda")
    X_.requires_grad_()

    # pyre-ignore[16]: TODO(T101400990): Pyre did not recognize the
    # `FusedLinear` import.
    triton_fused_linear = FusedLinear(
        shape[-1], shape[-1] // 2, bias=bias, activation=activation
    ).to("cuda")

    # Now check parity
    torch_linear.train()
    triton_fused_linear.train()

    torch_linear.zero_grad()
    triton_fused_linear.zero_grad()

    torch.testing.assert_close(
        triton_fused_linear.weight,
        torch_linear.weight,
        atol=1e-3,
        rtol=1e-3,
        msg="Broken test setup",
    )
    torch.testing.assert_close(X, X_, atol=1e-3, rtol=1e-3, msg="Broken test setup")

    with autocast(enabled=amp):
        y_torch = torch_sequence(X)
        y_triton = triton_fused_linear(X_)

        grad = torch.randn_like(y_torch)

        # Check that BW also gives the same result
        y_torch.backward(grad)
        y_triton.backward(grad)

        torch.testing.assert_close(X, X_, atol=1e-3, rtol=1e-3)

        # Input grad being correct checks both the loss + some of the backward pass
        assert X.grad is not None and X_.grad is not None
        torch.testing.assert_close(X.grad, X_.grad, atol=1e-3, rtol=1e-3)

        # Check that the linear layer bias are also properly trainable
        if bias:
            assert (
                triton_fused_linear.bias is not None
                and triton_fused_linear.bias.grad is not None
            )
            assert torch_linear.bias is not None and torch_linear.bias.grad is not None
            torch.testing.assert_close(
                torch_linear.bias.grad,
                triton_fused_linear.bias.grad,
                atol=1e-3,
                rtol=1e-3,
            )

        # Check that the linear layer weights are also properly trainable
        assert (
            torch_linear.weight.grad is not None
            and triton_fused_linear.weight.grad is not None
        )
        torch.testing.assert_close(
            torch_linear.weight.grad,
            triton_fused_linear.weight.grad,
            atol=1e-3,
            rtol=1e-3,
        )
//  [python]
def test_unrolled_cat(self):
        with inline_fusion_groups():
            def eager(x):
                ret = torch.empty(0)
                for i in range(x.shape[0]):
                    ret = torch.cat([ret, x[i].relu()])
                return ret
            script = torch.jit.script(eager)

            # Warm up with size=1 tensor; since the loop iterates once the
            # profile data will be "burned in" assuming size=1, and then
            # unrolled.
            x = torch.ones(1, 1)
            for _ in range(3):
                script(x)

            self.assertEqual(eager(x), script(x))

            # Now when an input hits the unrolled path, it will produce an
            # incorrectly-sized tensor, since size=1 has been burned in.
            x = torch.ones((8, 1))
            self.assertEqual(eager(x), script(x))
//  [python]
def check_jacrev_unary(group_fns, dim, batch_size, name):
    assert name in ["exp", "inv"]
    if not hasattr(torch, "vmap"):
        return

    test_fn = getattr(group_fns, name)

    fn_input = (
        torch.randn(batch_size, dim) if name == "exp" else group_fns.rand(batch_size)
    )

    def f(t):
        return group_fns.log(test_fn(t.unsqueeze(0))).squeeze(0)

    j = torch.vmap(torch.func.jacrev(f))(fn_input)
    jac_vmap = group_fns.left_project(fn_input, j) if name != "exp" else j

    jlog, jtest = [], []
    group_fns.log(test_fn(fn_input, jacobians=jtest), jacobians=jlog)
    jac_analytic = jlog[0] @ jtest[0]

    torch.testing.assert_allclose(jac_vmap, jac_analytic)
//  [python]
def check_jacrev_unary(group_fns, dim, batch_size, name):
    assert name in ["exp", "inv"]
    if not hasattr(torch, "vmap"):
        return

    test_fn = getattr(group_fns, name)

    fn_input = (
        torch.randn(batch_size, dim) if name == "exp" else group_fns.rand(batch_size)
    )

    def f(t):
        return group_fns.log(test_fn(t.unsqueeze(0))).squeeze(0)

    j = torch.vmap(torch.func.jacrev(f))(fn_input)
    jac_vmap = group_fns.left_project(fn_input, j) if name != "exp" else j

    jlog, jtest = [], []
    group_fns.log(test_fn(fn_input, jacobians=jtest), jacobians=jlog)
    jac_analytic = jlog[0] @ jtest[0]

    torch.testing.assert_close(jac_vmap, jac_analytic)
//  [python]
def test_batch_norm(self):
        def test(fn, args):
            trace = torch.jit.trace(fn, args)
            self.assertAllFused(trace.graph_for(*args))
            # TODO: Are `NaN`'s actually ok here or did this pass silently before, because `equal_nan=True` was the
            #  default?
            torch.testing.assert_allclose(fn(*args), trace(*args), equal_nan=True)

        def bn(i, x):
            return torch.batch_norm(i, x, x, x, x, False, 0.1, 1e-4, False).relu()

        def bn_no_weight(i, x):
            return torch.batch_norm(i, None, x, x, x, False, 0.1, 1e-4, False).relu()

        def bn_no_bias(i, x):
            return torch.batch_norm(i, x, None, x, x, False, 0.1, 1e-4, False).relu()

        def bn_neither(i, x):
            return torch.batch_norm(i, None, None, x, x, False, 0.1, 1e-4, False).relu()

        for device in self.devices:
            i = torch.randn(4, 16, 32, 40, device=device)
            x = torch.randn(16, device=device)
            for fn in [bn, bn_no_weight, bn_no_bias, bn_neither]:
                test(fn, (i, x))
//  [python]
def test_batch_norm(self):
        def test(fn, args):
            trace = torch.jit.trace(fn, args)
            self.assertAllFused(trace.graph_for(*args))
            # TODO: Are `NaN`'s actually ok here or did this pass silently before, because `equal_nan=True` was the
            #  default?
            torch.testing.assert_close(fn(*args), trace(*args), equal_nan=True)

        def bn(i, x):
            return torch.batch_norm(i, x, x, x, x, False, 0.1, 1e-4, False).relu()

        def bn_no_weight(i, x):
            return torch.batch_norm(i, None, x, x, x, False, 0.1, 1e-4, False).relu()

        def bn_no_bias(i, x):
            return torch.batch_norm(i, x, None, x, x, False, 0.1, 1e-4, False).relu()

        def bn_neither(i, x):
            return torch.batch_norm(i, None, None, x, x, False, 0.1, 1e-4, False).relu()

        for device in self.devices:
            i = torch.randn(4, 16, 32, 40, device=device)
            x = torch.randn(16, device=device)
            for fn in [bn, bn_no_weight, bn_no_bias, bn_neither]:
                test(fn, (i, x))
//  [python]
def test_one_hot(self, targets):
        composer_one_hot = _one_hot(targets)
        pytorch_one_hot = F.one_hot(targets)
        torch.testing.assert_allclose(composer_one_hot, pytorch_one_hot)
//  [python]
def test_one_hot(self, targets):
        composer_one_hot = _one_hot(targets)
        pytorch_one_hot = F.one_hot(targets)
        torch.testing.assert_close(composer_one_hot, pytorch_one_hot)
//  [python]
def check_log_map_passt(lie_group_fns, impl_module):
    set_global_params({"_allow_passthrough_ops": True})
    group = lie_group_fns.rand(
        4, device="cuda:0" if torch.cuda.is_available() else "cpu", requires_grad=True
    )
    jlist = []
    log_map_pt = lie_group_fns.log(group, jacobians=jlist)
    grad_pt = torch.autograd.grad(log_map_pt.sum(), group)
    log_map_ref = impl_module._log_autograd_fn(group)
    jac_ref = impl_module._jlog_impl(group)[0][0]
    grad_ref = torch.autograd.grad(log_map_ref.sum(), group)
    torch.testing.assert_allclose(log_map_pt, log_map_ref)
    torch.testing.assert_allclose(jlist[0], jac_ref)
    torch.testing.assert_allclose(grad_pt, grad_ref)
//  [python]
def check_log_map_passt(lie_group_fns, impl_module):
    set_global_params({"_allow_passthrough_ops": True})
    group = lie_group_fns.rand(
        4, device="cuda:0" if torch.cuda.is_available() else "cpu", requires_grad=True
    )
    jlist = []
    log_map_pt = lie_group_fns.log(group, jacobians=jlist)
    grad_pt = torch.autograd.grad(log_map_pt.sum(), group)
    log_map_ref = impl_module._log_autograd_fn(group)
    jac_ref = impl_module._jlog_impl(group)[0][0]
    grad_ref = torch.autograd.grad(log_map_ref.sum(), group)
    torch.testing.assert_close(log_map_pt, log_map_ref)
    torch.testing.assert_close(jlist[0], jac_ref)
    torch.testing.assert_close(grad_pt, grad_ref)
//  [python]
def check_lie_group_function(
    module, op_name: str, atol: float, inputs, funcs=None, batch_size=None
):
    op_impl = getattr(module, f"_{op_name}_impl")
    op = getattr(module, f"_{op_name}_autograd_fn")
    jop = getattr(module, f"_j{op_name}_autograd_fn")

    # Check jacobians
    jacs_impl = torch.autograd.functional.jacobian(op_impl, inputs, vectorize=True)
    jacs = torch.autograd.functional.jacobian(op, inputs, vectorize=True)

    if funcs is None:
        for jac_impl, jac in zip(jacs_impl, jacs):
            torch.testing.assert_allclose(jac_impl, jac, atol=atol, rtol=atol)
    else:
        for jac_impl, jac, func in zip(jacs_impl, jacs, funcs):
            if func is None:
                torch.testing.assert_allclose(jac_impl, jac, atol=atol, rtol=atol)
            else:
                torch.testing.assert_allclose(
                    func(jac_impl), func(jac), atol=atol, rtol=atol
                )

    # Check multi-batch consistency
    if batch_size is None:
        return
    lb = len(batch_size)
    flattened_inputs = [x.reshape(-1, *x.shape[lb:]) for x in inputs]
    out = op(*inputs)
    flattened_out = op(*flattened_inputs)
    if jop is None:
        return
    jout = jop(*inputs)[0]
    flattened_jout = jop(*flattened_inputs)[0]
    torch.testing.assert_allclose(out, flattened_out.reshape(*batch_size, *out.shape[lb:]))
    for j, jf in zip(jout, flattened_jout):
        torch.testing.assert_allclose(j, jf.reshape(*batch_size, *j.shape[lb:]))