test_tensorexpr.py

# Owner(s): ["NNC"]
# ruff: noqa: F841

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
import unittest
import itertools

from torch.testing._internal.common_utils import suppress_warnings, num_profiled_runs, run_tests, skipIfTorchDynamo

from torch.testing._internal.jit_utils import JitTestCase, TensorExprTestOptions

LLVM_ENABLED = torch._C._llvm_enabled()

class BaseTestClass(JitTestCase):
    def setUp(self):
        super().setUp()
        self.tensorexpr_options = TensorExprTestOptions()
        self.devices = ['cpu'] if not torch.cuda.is_available() else ['cpu', 'cuda']
        self.dtypes = [torch.float32, torch.bfloat16] if LLVM_ENABLED else [torch.float32]

    def tearDown(self):
        self.tensorexpr_options.restore()
        super().tearDown()

    def assertLastGraphAllFused(self):
        self.assertAllFused(torch.jit.last_executed_optimized_graph())


def warmup_and_run_forward(f, *args):
    for _ in range(torch._C._jit_get_num_profiled_runs() + 1):
        results = f(*args)
    return results


@skipIfTorchDynamo()
class TestTensorExprFuser(BaseTestClass):
    def test_easy(self):
        def easy(x, y):
            aaa = torch.add(x, y)
            return aaa

        traced = torch.jit.trace(easy, (torch.rand(1024), torch.rand(1024)))

        a = torch.rand(1024)
        b = torch.rand(1024)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(a.numpy() + b.numpy(), x.numpy())

    def test_three_arg(self):
        def easy(x, y, z):
            aaa = torch.add(x, y)
            bbb = torch.add(aaa, z)
            return bbb

        traced = torch.jit.trace(
            easy, (torch.rand(1024), torch.rand(1024), torch.rand(1024))
        )

        a = torch.rand(1024)
        b = torch.rand(1024)
        c = torch.rand(1024)
        x = warmup_and_run_forward(traced, a, b, c)
        self.assertLastGraphAllFused()
        npr = a.numpy() + b.numpy() + c.numpy()
        np.testing.assert_allclose(npr, x.numpy())

    def test_four_arg(self):
        def run_addcmul(x, y, z, w):
            c = torch.addcmul(torch.add(x, y), z, w)
            return c

        for dev in self.devices:
            rand_a = torch.rand(1024, dtype=torch.float, device=dev)
            rand_b = torch.rand(1024, dtype=torch.float, device=dev)
            rand_c = torch.rand(1024, dtype=torch.float, device=dev)
            rand_d = torch.rand(1024, dtype=torch.float, device=dev)

            traced = torch.jit.trace(
                run_addcmul,
                (
                    torch.zeros(1024, dtype=torch.float, device=dev),
                    torch.zeros(1024, dtype=torch.float, device=dev),
                    torch.zeros(1024, dtype=torch.float, device=dev),
                    torch.zeros(1024, dtype=torch.float, device=dev),
                ),
            )

            x = warmup_and_run_forward(traced, rand_a, rand_b, rand_c, rand_d)
            self.assertLastGraphAllFused()
            y = run_addcmul(rand_a, rand_b, rand_c, rand_d)
            np.testing.assert_allclose(x.cpu().numpy(), y.cpu().numpy(), atol=1e-6)

    def test_three_arg2(self):
        for device in self.devices:
            def test(x, y, z):
                aaa = torch.add(x, y)
                bbb = torch.add(aaa, z)
                return bbb

            M = 32
            N = 32
            traced = torch.jit.trace(
                test,
                (
                    torch.rand(M, N, device=device),
                    torch.rand(M, N, device=device),
                    torch.rand(M, N, device=device),
                ),
            )

            a = torch.rand(M, N, device=device)
            b = torch.rand(M, N, device=device)
            c = torch.rand(M, N, device=device)
            x = traced(a, b, c)
            x = warmup_and_run_forward(traced, a, b, c)
            self.assertLastGraphAllFused()
            npr = a.cpu().numpy() + b.cpu().numpy() + c.cpu().numpy()
            np.testing.assert_allclose(npr, x.cpu().numpy())

    def test_broadcast3(self):
        for device in self.devices:
            def test_body(M, N, L, K):
                def test(x, y, z):
                    v1 = torch.add(x, y)
                    v2 = torch.add(v1, z)
                    return v2

                a_shape = [M, N]
                b_shape = [L, M, 1]
                c_shape = [K, L, 1, 1]
                traced = torch.jit.trace(
                    test,
                    (
                        torch.rand(*a_shape, device=device),
                        torch.rand(*b_shape, device=device),
                        torch.rand(*c_shape, device=device),
                    ),
                )

                a = torch.rand(*a_shape, device=device)
                b = torch.rand(*b_shape, device=device)
                c = torch.rand(*c_shape, device=device)
                x = warmup_and_run_forward(traced, a, b, c)
                self.assertLastGraphAllFused()
                npr = a.cpu().numpy() + b.cpu().numpy() + c.cpu().numpy()
                np.testing.assert_allclose(npr, x.cpu().numpy())

            test_configs = [[5, 2, 7, 3], [8, 8, 8, 8]]
            for test_config in test_configs:
                test_body(*test_config)

    def test_all_combos(self):
        def easy(x, y, z):
            a = torch.add(x, y)
            b = torch.add(a, z)
            c = torch.add(x, b)
            d = torch.add(c, a)
            return d

        def np_easy(x, y, z):
            a = x + y
            b = a + z
            c = x + b
            d = c + a
            return d

        traced = torch.jit.trace(
            easy, (torch.rand(1024), torch.rand(1024), torch.rand(1024))
        )

        a = torch.rand(1024)
        b = torch.rand(1024)
        c = torch.rand(1024)
        x = warmup_and_run_forward(traced, a, b, c)
        self.assertLastGraphAllFused()
        npr = np_easy(a.numpy(), b.numpy(), c.numpy())
        np.testing.assert_allclose(npr, x.numpy())

    def test_rank_two(self):
        def easy(x, y, z):
            a = torch.add(x, y)
            b = torch.add(a, z)
            c = torch.add(x, b)
            d = torch.add(c, a)
            return d

        def np_easy(x, y, z):
            a = x + y
            b = a + z
            c = x + b
            d = c + a
            return d

        shape = 32, 32
        traced = torch.jit.trace(
            easy, (torch.rand(shape), torch.rand(shape), torch.rand(shape))
        )

        a = torch.rand(shape)
        b = torch.rand(shape)
        c = torch.rand(shape)
        x = warmup_and_run_forward(traced, a, b, c)
        self.assertLastGraphAllFused()
        npr = np_easy(a.numpy(), b.numpy(), c.numpy())
        np.testing.assert_allclose(npr, x.numpy())

    def test_broadcast(self):
        def easy(x, y, z):
            a = torch.add(x, y)
            b = torch.add(a, z)
            return b

        def np_easy(x, y, z):
            a = x + y
            b = a + z
            return b

        N = 32
        traced = torch.jit.trace(easy, (torch.rand(N, N), torch.rand(N), torch.rand(N, N)))

        a = torch.rand(N, N)
        b = torch.rand(N)
        c = torch.rand(N, N)
        x = warmup_and_run_forward(traced, a, b, c)
        self.assertLastGraphAllFused()
        npr = np_easy(a.numpy(), b.numpy(), c.numpy())
        np.testing.assert_allclose(npr, x.numpy())

    def test_broadcast_2(self):
        zero = torch.tensor([0.0], dtype=torch.float)

        def foo(x, y, z):
            aaa = torch.add(x, y)
            bbb = torch.add(zero, aaa)
            return torch.add(bbb, z)

        def foo_np(x, y, z):
            a = x + y
            b = zero.numpy() + a
            return b + z

        x = torch.rand(3, 4)
        y = torch.ones(3, 1)
        z = torch.rand(4)
        traced = torch.jit.trace(foo, (x, y, z))

        r = warmup_and_run_forward(traced, x, y, z)
        self.assertLastGraphAllFused()

        rnp = foo_np(x.numpy(), y.numpy(), z.numpy())
        np.testing.assert_allclose(r, rnp)

    def test_broadcast_big2(self):
        zero = torch.tensor([0.0], dtype=torch.float)

        def foo(x, y, z):
            aaa = torch.add(x, y)
            bbb = torch.add(zero, aaa)
            return torch.add(bbb, z)

        def foo_np(x, y, z):
            a = x + y
            b = zero.numpy() + a
            return b + z

        x = torch.rand(32, 1024)
        y = torch.ones(32, 1)
        z = torch.rand(1024)
        traced = torch.jit.trace(foo, (x, y, z))

        r = warmup_and_run_forward(traced, x, y, z)
        self.assertLastGraphAllFused()
        rnp = foo_np(x.numpy(), y.numpy(), z.numpy())
        np.testing.assert_allclose(r, rnp)

    def test_alpha(self):
        def alpha(x):
            aaa = torch.add(x, x, alpha=2.0)
            return aaa

        traced = torch.jit.trace(alpha, (torch.tensor([1.0])))

        a = torch.tensor([1.0])
        x = traced(a)
        np.testing.assert_allclose(a.numpy() + 2.0 * a.numpy(), x.numpy())

    @suppress_warnings
    def test_constant(self):
        def constant(x):
            bbb = torch.tensor([1.0])
            aaa = torch.add(x, bbb)
            return aaa

        traced = torch.jit.trace(constant, (torch.tensor([1.0])))

        a = torch.tensor([1.0])
        x = warmup_and_run_forward(traced, a)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(a.numpy() + 1.0, x.numpy())

    def test_add_sub(self):
        def easy(x, y, z):
            aaa = torch.add(x, y)
            bbb = torch.sub(aaa, z)
            return bbb

        traced = torch.jit.trace(
            easy, (torch.rand(1024), torch.rand(1024), torch.rand(1024))
        )

        a = torch.rand(1024)
        b = torch.rand(1024)
        c = torch.rand(1024)
        x = warmup_and_run_forward(traced, a, b, c)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(a.numpy() + b.numpy() - c.numpy(), x.numpy())

    def test_promotion(self):
        def easy(x, y):
            aaa = torch.add(x, y)
            return aaa

        traced = torch.jit.trace(
            easy,
            (torch.zeros(1024, dtype=torch.int32), torch.rand(1024, dtype=torch.float32)),
        )

        a = torch.zeros(1024, dtype=torch.int32)
        b = torch.rand(1024, dtype=torch.float32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(a.numpy() + b.numpy(), x.numpy())

    def test_double(self):
        TENSOR_LEN = 8

        def easy(x, y):
            aaa = torch.add(x, y)
            bbb = torch.mul(aaa, y)
            return bbb

        traced = torch.jit.trace(
            easy,
            (torch.rand(TENSOR_LEN, dtype=torch.float64), torch.full((TENSOR_LEN,), 0.5, dtype=torch.float64)),
        )

        a = torch.rand(TENSOR_LEN, dtype=torch.double)
        b = torch.full((TENSOR_LEN,), 0.5, dtype=torch.double)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose((a.numpy() + b.numpy()) * b.numpy(), x.numpy())

    def test_short(self):
        TENSOR_LEN = 8

        def easy(x, y):
            aaa = torch.add(x, y)
            bbb = torch.mul(aaa, y)
            return bbb

        traced = torch.jit.trace(
            easy,
            (torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int16),
             torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int16)),
        )

        a = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int16)
        b = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int16)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose((a.numpy() + b.numpy()) * b.numpy(), x.numpy())

    def test_char(self):
        TENSOR_LEN = 8

        def easy(x, y):
            aaa = torch.add(x, y)
            bbb = torch.mul(aaa, y)
            return bbb

        traced = torch.jit.trace(
            easy,
            (torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8),
             torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8)),
        )

        a = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8)
        b = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose((a.numpy() + b.numpy()) * b.numpy(), x.numpy())

    def test_int64_promotion(self):
        TENSOR_LEN = 8

        def easy(x, y):
            aaa = torch.add(x, y)
            bbb = torch.mul(aaa, y)
            return bbb

        traced = torch.jit.trace(
            easy,
            (torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8),
             torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int64)),
        )

        a = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int8)
        b = torch.randint(TENSOR_LEN, (TENSOR_LEN,), dtype=torch.int64)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose((a.numpy() + b.numpy()) * b.numpy(), x.numpy())

    def test_eq(self):
        def easy(x, y):
            c = torch.eq(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.zeros(1024), torch.zeros(1024)))
        a = torch.zeros(1024, dtype=torch.int32)
        b = torch.zeros(1024, dtype=torch.int32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(np.ones(1024), x.numpy())

    def test_ne(self):
        def easy(x, y):
            c = torch.ne(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.zeros(1024), torch.zeros(1024)))
        a = torch.zeros(1024, dtype=torch.int32)
        b = torch.ones(1024, dtype=torch.int32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(np.ones(1024), x.numpy())

    def test_ge(self):
        def easy(x, y):
            c = torch.ge(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.zeros(1024), torch.zeros(1024)))
        aa = np.empty([1024], dtype=np.int32)
        aa.fill(5)
        a = torch.from_numpy(aa)
        b = torch.zeros(1024, dtype=torch.int32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(np.ones(1024), x.numpy())

    def test_gt(self):
        def easy(x, y):
            c = torch.gt(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.zeros(1024), torch.zeros(1024)))
        a = torch.ones(1024, dtype=torch.int32)
        b = torch.zeros(1024, dtype=torch.int32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(np.ones(1024), x.numpy())

    def test_le(self):
        def easy(x, y):
            c = torch.le(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.zeros(1024), torch.zeros(1024)))
        aa = np.empty([1024], dtype=np.int32)
        aa.fill(5)
        a = torch.from_numpy(aa)
        b = torch.zeros(1024, dtype=torch.int32)
        x = warmup_and_run_forward(traced, a, b)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(np.zeros(1024), x.numpy())

    def test_lt(self):
        def easy(x, y):
            c = torch.lt(x, y)
            return c

        for dev in self.devices:
            traced = torch.jit.trace(easy, (torch.zeros(1024, device=dev), torch.zeros(1024, device=dev)))
            a = torch.ones(1024, dtype=torch.int32, device=dev)
            b = torch.zeros(1024, dtype=torch.int32, device=dev)
            x = warmup_and_run_forward(traced, a, b)
            self.assertLastGraphAllFused()
            np.testing.assert_allclose(np.zeros(1024), x.cpu().numpy())

    @suppress_warnings
    def test_min_max(self):
        def test(x, y):
            return torch.max(torch.min(x, y), torch.tensor([4.0]))

        traced = torch.jit.trace(test, (torch.zeros(1024), torch.zeros(1024)))
        a = 8.0 * torch.rand(1024)
        b = 8.0 * torch.rand(1024)
        np.testing.assert_allclose(
            warmup_and_run_forward(traced, a, b), np.maximum(np.minimum(a.numpy(), b.numpy()), [4.0])
        )
        self.assertLastGraphAllFused()

    def test_min_max_reduction(self):
        def test(x):
            return torch.min(x) + torch.max(x)

        traced = torch.jit.trace(test, (torch.zeros(1024)))
        a = 8.0 * torch.rand(1024)
        np.testing.assert_allclose(warmup_and_run_forward(traced, a), np.amin(a.numpy()) + np.amax(a.numpy()))
        self.assertLastGraphAllFused()

    def test_min_max_reduction2(self):
        def test(x):
            return x.min() + x.max()

        traced = torch.jit.trace(test, (torch.zeros(1024)))
        a = 8.0 * torch.rand(1024)
        np.testing.assert_allclose(warmup_and_run_forward(traced, a), np.amin(a.numpy()) + np.amax(a.numpy()))
        self.assertLastGraphAllFused()

    def test_min_max_reduction_dim1(self):
        def test(x):
            return torch.min(x, 1)[0] + torch.max(x, 1)[0]

        traced = torch.jit.trace(test, (torch.zeros(16, 16)))
        a = 8.0 * torch.rand(16, 16)
        np.testing.assert_allclose(warmup_and_run_forward(traced, a), np.amin(
            a.numpy(), axis=1) + np.amax(a.numpy(), axis=1))
        self.assertLastGraphAllFused()

    def test_min_max_reduction_dim1_2(self):
        def test(x):
            return torch.min(x * x, 1)

        traced = torch.jit.trace(test, (torch.zeros(16, 16)))
        a = 8.0 * torch.rand(16, 16)
        np.testing.assert_allclose(warmup_and_run_forward(traced, a)[0], np.amin((a * a).numpy(), axis=1))
        self.assertLastGraphAllFused()

    def test_clamp(self):
        def test(x):
            return torch.clamp(x + 3.0, 0.0, 6.0)

        for dev in self.devices:
            traced = torch.jit.trace(test, (torch.zeros(1024, device=dev)))
            a = 20.0 * torch.rand(1024, device=dev) - 10.0
            an = a.cpu().numpy()
            np.testing.assert_allclose(warmup_and_run_forward(traced, a).cpu(), np.clip(an + 3.0, 0.0, 6.0))
            self.assertLastGraphAllFused()

    def test_relu(self):
        def test(x):
            return torch.clamp(F.relu(x), 0, 0.5)

        for dev in self.devices:
            traced = torch.jit.trace(test, (torch.zeros(1024, device=dev)))
            a = 20.0 * torch.rand(1024, device=dev) - 10.0
            an = a.cpu().numpy()
            np.testing.assert_allclose(warmup_and_run_forward(traced, a).cpu(), np.clip((np.maximum(0, an)), 0, 0.5))
            self.assertLastGraphAllFused()

    def test_reps(self):
        def easy(x, y):
            c = torch.add(x, y)
            return c

        traced = torch.jit.trace(easy, (torch.rand(1024), torch.rand(1024)))

        for _ in range(32):
            a = torch.ones(1024)
            b = torch.zeros(1024)
            x = warmup_and_run_forward(traced, a, b)
            np.testing.assert_allclose(np.ones(1024), x.numpy())

    def test_add_const_rhs(self):
        def test(x):
            return x + 3.0

        traced = torch.jit.trace(test, torch.rand(4))
        x = torch.rand(4)
        y = warmup_and_run_forward(traced, x)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(x.numpy() + 3.0, y.numpy())

    def test_int_output(self):
        def test(x, y, z):
            return x * y * z

        xs = [(torch.rand(4) * 3 + 1).to(torch.int32) for i in range(3)]
        x, y, z = xs
        xn, yn, zn = (t.numpy() for t in xs)
        traced = torch.jit.trace(test, (x, y, z))
        res = warmup_and_run_forward(traced, x, y, z)
        self.assertLastGraphAllFused()
        np.testing.assert_allclose(xn * yn * zn, res.numpy())

    def test_binary_ops(self):
        def test_atan2(x, y):
            c = torch.atan2(torch.add(x, y), y)
            return c

        def test_gt(x, y):
            c = torch.gt(torch.add(x, y), y)
            return c

        def test_ge(x, y):
            c = torch.ge(torch.add(x, y), y)
            return c

        def test_lt(x, y):
            c = torch.lt(torch.add(x, y), y)
            return c

        def test_le(x, y):
            c = torch.le(torch.add(x, y), y)
            return c

        def test_lerp(x, y):
            c = torch.lerp(torch.add(x, 1), x, 2.0)
            return c

        def test_mul(x, y):
            c = torch.mul(torch.add(x, y), y)
            return c

        def test_ne(x, y):
            c = torch.ne(torch.add(x, y), y)
            return c

        def test_div(x, y):
            c = torch.div(torch.add(x, y), 2)
            return c

        def test_eq(x, y):
            c = torch.eq(torch.add(x, y), y)
            return c

        def test_fmod(x, y):
            c = torch.fmod(torch.add(x, y), 2)
            return c

        def test_sub(x, y):
            c = torch.sub(torch.add(x, y), x)
            return c

        def test_remainder(x, y):
            c = torch.remainder(torch.add(x, y), 3.0)
            return c

        def test_pow(x, y):
            c = torch.pow(torch.add(x, y), 2.0)
            return c

        def test_type_as(x, y):
            return x.type_as(torch.add(x, y))

        cmp_fns = {
            test_gt,
            test_ge,
            test_lt,
            test_le,
            test_ne,
            test_eq
        }

        non_cmp_fns = {
            test_atan2,
            test_lerp,
            test_mul,
            test_div,
            test_fmod,
            test_sub,
            test_remainder,
            test_pow,
            test_type_as,
        }

        all_test_fns = cmp_fns.union(non_cmp_fns)
        fn_dev_dtype = itertools.product(all_test_fns, self.devices, self.dtypes)
        for torch_fn, dev, data_type in fn_dev_dtype:
            if torch_fn is test_lerp and data_type is torch.bfloat16:
                continue
            rand_a = torch.rand(1024, dtype=data_type, device=dev)
            rand_b = torch.rand(1024, dtype=data_type, device=dev)
            in1 = 20 * torch.rand(1024, dtype=data_type, device=dev)
            in2 = 20 * torch.rand(1024, dtype=data_type, device=dev)
            traced = torch.jit.trace(torch_fn, (in1, in2))
            x = warmup_and_run_forward(traced, rand_a, rand_b)
            self.assertLastGraphAllFused()

            _atol = 2e-3
            _rtol = 1e-5
            if data_type is torch.bfloat16:
                # Compared to aten logic, NNC coudl save addtional BF16/Fp32 conversion.
                # Take d = a + b - c as an example, the aten logic is as follows at
                # operator level:
                #    tmp = to_bf16(to_fp32(a) + to_fp32(b))
                #    d = to_bf16(to_fp32(tmp) + to_fp32(c))
                # But NNC could fuse the compression and remove the redudant conversions.
                # The final statement is as follows
                #    d = to_bf16(to_fp32(a) + to_fp32(b) + to_fp32(c))
                # Hence, we simulate NNC computation by feeding fp32 tensors and converting
                # the result tensor back to bf16. The simulation could avoid the numeric
                # deviation to simplify the result comprasion
                y = warmup_and_run_forward(traced, rand_a.float(), rand_b.float())
                if torch_fn not in cmp_fns:
                    y = y.bfloat16()
                _atol = 2e-2
            else:
                y = torch_fn(rand_a, rand_b)
            self.assertEqual(x.cpu(), y.cpu(), atol=_atol, rtol=_rtol)

    def test_unary_ops(self):
        def test_cast_float(x, y):
            c = torch.ops.aten._cast_Float(torch.add(x, y))
            return c

        def test_round(x, y):
            c = torch.round(torch.add(x, y))
            return c

        def test_sin(x, y):
            c = torch.sin(torch.add(x, y))
            return c

        def test_asin(x, y):
            c = torch.asin(torch.add(x, y))
            return c

        def test_sinh(x, y):
            c = torch.sinh(torch.add(x, y))
            return c

        def test_cos(x, y):
            c = torch.cos(torch.add(x, y))
            return c

        def test_acos(x, y):
            c = torch.acos(torch.add(x, y))
            return c

        def test_cosh(x, y):
            c = torch.cosh(torch.add(x, y))
            return c

        def test_tan(x, y):
            c = torch.tan(torch.add(x, y))
            return c

        def test_atan(x, y):
            c = torch.atan(torch.add(x, y))
            return c

        def test_tanh(x, y):
            c = torch.tanh(torch.add(x, y))
            return c

        def test_sqrt(x, y):
            c = torch.sqrt(torch.add(x, y))
            return c

        def test_rsqrt(x, y):
            c = torch.rsqrt(torch.add(x, y))
            return c

        def test_floor(x, y):
            c = torch.floor(torch.add(x, y))
            return c

        def test_ceil(x, y):
            c = torch.ceil(torch.add(x, y))
            return c

        def test_trunc(x, y):
            c = torch.trunc(torch.add(x, y))
            return c

        def test_abs(x, y):
            c = torch.abs(torch.add(x, y))
            return c

        def test_log(x, y):
            c = torch.log(torch.add(x, y))
            return c

        def test_log2(x, y):
            c = torch.log2(torch.add(x, y))
            return c

        def test_log10(x, y):
            c = torch.log10(torch.add(x, y))
            return c

        def test_log1p(x, y):
            c = torch.log1p(torch.add(x, y))
            return c

        def test_rqrt(x, y):
            c = torch.rsqrt(torch.add(x, y))
            return c

        def test_erf(x, y):
            c = torch.erf(torch.add(x, y))
            return c

        def test_exp(x, y):
            c = torch.exp(torch.add(x, y))
            return c

        def test_expm1(x, y):
            c = torch.expm1(torch.add(x, y))
            return c

        def test_erfc(x, y):
            c = torch.erfc(torch.add(x, y))
            return c

        def test_frac(x, y):
            c = torch.frac(torch.add(x, y))
            return c

        def test_lgamma(x, y):
            c = torch.lgamma(torch.add(x, y))
            return c

        def test_sigmoid(x, y):
            c = torch.sigmoid(torch.add(x, y))
            return c

        def test_reciprocal(x, y):
            c = torch.reciprocal(torch.add(x, y))
            return c

        def test_neg(x, y):
            c = torch.neg(torch.add(x, y))
            return c

        def test_relu(x, y):
            c = torch.relu(torch.add(x, y))
            return c

        def test_hardtanh(x, y):
            c = F.hardtanh(torch.add(x, y), -1.0, 1.0)
            return c

        def test_threshold(x, y):
            c = F.threshold(torch.add(x, y), 0.5, 10)
            return c

        gpu_only_fns = {
            test_erf,
            test_erfc
        }
        fns = {
            test_round,
            test_sin,
            test_asin,
            test_sinh,
            test_cos,
            test_acos,
            test_cosh,
            test_tan,
            test_atan,
            test_sqrt,
            test_floor,
            test_ceil,
            test_trunc,
            test_abs,
            test_log,
            test_log2,
            test_log10,
            test_log1p,
            test_rsqrt,
            test_exp,
            test_expm1,
            test_frac,
            test_lgamma,
            test_reciprocal,
            test_neg,
            test_threshold,
            test_relu,
            test_tanh,
            test_hardtanh,
            test_sigmoid,
        }
        fn_dev_dtype = itertools.product(gpu_only_fns.union(fns), self.devices, self.dtypes)

        torch.manual_seed(0)
        for torch_fn, dev, data_type in fn_dev_dtype:
            if torch_fn == test_lgamma and dev == "cuda":
                # lgamma_cuda does not support BF16
                continue
            rand_a = torch.rand(1024, dtype=data_type, device=dev)
            rand_b = torch.rand(1024, dtype=data_type, device=dev)

            ins = 20 * torch.rand(1024, dtype=data_type, device=dev)
            cc = np.empty([1024], dtype=np.float32)
            cc.fill(np.nan)
            nans = torch.from_numpy(cc).to(dev)
            traced = torch.jit.trace(torch_fn, (ins, ins))
            x = warmup_and_run_forward(traced, rand_a, rand_b)
            self.assertLastGraphAllFused()

            _atol = 5e-3 if data_type is torch.bfloat16 else 2e-3
            _rtol = 1e-5
            if data_type is torch.bfloat16 and torch_fn not in gpu_only_fns:
                y = warmup_and_run_forward(traced, rand_a.float(), rand_b.float())
                y = y.bfloat16()
            else:
                y = torch_fn(rand_a, rand_b)

            self.assertEqual(x.cpu(), y.cpu(), atol=_atol, rtol=_rtol)
            # nans
            # TODO: reenable. Currently all of the tests fail
            # traced = torch.jit.trace(torch_fn, (ins, ins))
            # x = warmup_and_run_forward(traced, rand_a, rand_b)
            # y = torch_fn(nans, rand_b)
            # try:
            #     np.testing.assert_allclose(x.cpu().numpy(), y.cpu().numpy())
            #     print("Succeeded on dev=", dev, "function=", torch_fn)
            # except AssertionError:
            #     # Print extra info before exiting:
            #     print("Failed on dev=", dev, "function=", torch_fn)
            #     # np.testing.assert_allclose(x.cpu().numpy(), y.cpu().numpy())


    def test_round_2(self):
        def round(x):
            return torch.round(x)

        for data_type in [torch.float32, torch.double]:
            a = torch.tensor([0.2, 1.6, 2.5, 3.5]).to(data_type)
            traced = torch.jit.trace(round, (a))
            x = warmup_and_run_forward(traced, a)
            self.assertLastGraphAllFused()
            y = round(x)
            self.assertEqual(x, y)

    def test_rand_like(self):
        N = 1 << 16

        def run_rand_like(x, y):
            return torch.rand_like(torch.add(x, y))

        for device in self.devices:
            x = torch.rand(N, device=device)
            traced = torch.jit.trace(run_rand_like, (x, x), check_trace=False)

            for data_type in self.dtypes:
                _x = x.to(dtype=data_type)
                x_v = warmup_and_run_forward(traced, _x, _x)
                self.assertLastGraphAllFused()

            x_np = x.cpu().numpy()
            x1_mean = np.mean(x_np)
            x2_mean = np.mean(x_np ** 2)
            x3_mean = np.mean(x_np ** 3)
            np.testing.assert_allclose(x1_mean, 1. / 2, rtol=2e-2)
            np.testing.assert_allclose(x2_mean, 1. / 3, rtol=2e-2)
            np.testing.assert_allclose(x3_mean, 1. / 4, rtol=2e-2)

    def test_nans(self):
        def test_max(x, y):
            return torch.max(2 * x, 2 * y)

        def test_min(x, y):
            return torch.min(2 * x, 2 * y)

        tmax = torch.jit.trace(test_max, (torch.rand(1), torch.rand(1)))
        tmin = torch.jit.trace(test_min, (torch.rand(1), torch.rand(1)))

        for data_type in self.dtypes:
            x = torch.tensor([np.nan]).to(dtype=data_type)
            y = torch.tensor([1.0]).to(dtype=data_type)

        assert np.isnan(warmup_and_run_forward(tmin, x, y).float().item())
        assert np.isnan(warmup_and_run_forward(tmin, y, x).float().item())
        self.assertLastGraphAllFused()
        assert np.isnan(warmup_and_run_forward(tmax, x, y).float().item())
        assert np.isnan(warmup_and_run_forward(tmax, y, x).float().item())
        self.assertLastGraphAllFused()

    def test_double_intrinsics(self):
        def do_pow(x):
            return torch.pow(x, 7)

        for device in self.devices:
            x = torch.rand(10, dtype=torch.double, device=device)
            traced = torch.jit.trace(do_pow, (x))
            x = warmup_and_run_forward(traced, x)
            self.assertLastGraphAllFused()

    def test_remainder(self):
        def run_remainder(x, y):
            c = torch.remainder(torch.add(x, y), x)
            return c

        for data_type in self.dtypes:
            a = torch.rand(1024, dtype=data_type)
            b = torch.rand(1024, dtype=data_type)
            zeros = torch.zeros(1024, dtype=data_type)
            cc = np.array(1024, dtype=float)
            cc.fill(np.nan)
            nans = torch.from_numpy(cc).to(dtype=data_type)

            # random floats
            zeros1 = torch.zeros(1024, dtype=data_type)
            zeros2 = torch.zeros(1024, dtype=data_type)

            traced = torch.jit.trace(run_remainder, (zeros1, zeros2))
            x = warmup_and_run_forward(traced, a, b)
            self.assertLastGraphAllFused()
            y = run_remainder(a, b)
            if data_type is torch.bfloat16:
                self.assertEqual(x, y, atol=4e-3, rtol=2e-3)
            else:
                self.assertEqual(x, y)

            # div by 0
            traced = torch.jit.trace(run_remainder, (zeros1, zeros2))
            x = warmup_and_run_forward(traced, zeros, a)
            self.assertLastGraphAllFused()
            y = run_remainder(zeros, a)
            self.assertEqual(x, y)

            # numerators and denominatos are nan
            traced = torch.jit.trace(run_remainder, (zeros1, zeros2))
            x = warmup_and_run_forward(traced, nans, a)
            self.assertLastGraphAllFused()
            y = run_remainder(nans, a)
            self.assertEqual(x, y)

    def test_multioutput(self):
        def easy(x):
            b = x + 1
            c = b + b
            return (b, c)

        traced = torch.jit.trace(easy, (torch.zeros(1024)))

        a = torch.zeros(1024)
        b, c = warmup_and_run_forward(traced, a)
        self.assertLastGraphAllFused()
        bp = a.numpy() + 1
        cp = bp + bp
        np.testing.assert_allclose(b.numpy(), bp)
        np.testing.assert_allclose(c.numpy(), cp)

    def test_chunk(self):
        def easy(x):
            y = x + 1
            aaa, bbb = torch.chunk(y, 2)
            return aaa + bbb

        for data_type in self.dtypes:
            trace_input = torch.zeros(1024, 1024, dtype=data_type)
            traced = torch.jit.trace(easy, (trace_input))

            a = torch.zeros(32, 32, dtype=data_type)
            x = warmup_and_run_forward(traced, a)
            self.assertLastGraphAllFused()
            npr = a.float().numpy()
            npr2 = npr + 1
            npr_a, npr_b = np.array_split(npr2, 2)
            np.testing.assert_allclose(npr_a + npr_b, x.float().numpy())

    def test_cat(self):
        for device in self.devices:
            _dim = 1

            def foo(*args):
                args_2 = [v + i for i, v in enumerate(args)]
                v = torch.cat(args_2, dim=_dim)
                return v * v

            for data_type in self.dtypes:
                M = 16
                Ns = [128, 16, 1]
                values = [torch.zeros(M, N, dtype=data_type, device=device) for N in Ns]
                traced = torch.jit.trace(foo, values)

                x = warmup_and_run_forward(traced, *values)
                self.assertLastGraphAllFused()
                ref = foo(*values)
                np.testing.assert_allclose(ref.cpu().float().numpy(), x.cpu().float().numpy())

            # Test channels-last
            for _cur_dim in range(4):
                _dim = _cur_dim
                values = [torch.randn((2, 3, 4, 5), device=device).to(memory_format=torch.channels_last) for _ in range(10)]
                traced = torch.jit.trace(foo, values)

                x = warmup_and_run_forward(traced, *values)
                self.assertLastGraphAllFused()
                ref = foo(*values)
                self.assertEqual(ref, x)

    # This test checks that we correctly handle fusion group with just aten::cat in it.
    # Note that the test only makes sense with min_fusion_group=1, otherwise no
    # fusion groups would be formed at all.
    # TODO: Fix and re-enable the test.
    @unittest.skip("cat is broken with fusion group inlining disabled")
    def test_cat_only(self):
        for device in self.devices:
            def foo(*args):
                args_2 = [v + i for i, v in enumerate(args)]
                v = torch.cat(args_2, dim=1)
                return v

            M = 16
            Ns = [128, 16, 1]
            values = [torch.zeros(M, N, device=device) for N in Ns]
            traced = torch.jit.trace(foo, values)

            x = warmup_and_run_forward(traced, *values)
            self.assertLastGraphAllFused()
            ref = foo(*values)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

    def test_cat_negative_dim(self):
        for device in self.devices:
            def foo(*args):
                v = torch.cat(args, dim=-1)
                return v * v

            M = 16
            Ns = [128, 16, 1]
            values = [torch.randn(M, N, device=device) for N in Ns]
            traced = torch.jit.trace(foo, values)

            x = warmup_and_run_forward(traced, *values)
            self.assertLastGraphAllFused()
            ref = foo(*values)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

    def test_cat_promote_inputs(self):
        for device in self.devices:
            def foo(*args):
                v = torch.cat(args, dim=1)
                return v * v

            M = 16
            Ns = [128, 16, 1]
            dtypes = [torch.half, torch.float32, torch.double]
            values = [torch.randn(M, N, device=device, dtype=dt) for N, dt in zip(Ns, dtypes)]
            traced = torch.jit.trace(foo, values)

            x = warmup_and_run_forward(traced, *values)
            self.assertLastGraphAllFused()
            ref = foo(*values)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

    def test_cat_empty_tensors(self):
        for device in self.devices:
            def foo(*args):
                v = torch.cat(args, dim=1)
                return v * v

            M = 16
            Ns = [128, 16, 1]
            empty = torch.tensor([], device=device, dtype=torch.double)
            values = [empty] + [torch.randn(M, N, device=device) for N in Ns]
            traced = torch.jit.trace(foo, values)

            x = warmup_and_run_forward(traced, *values)
            self.assertLastGraphAllFused()
            ref = foo(*values)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

            # now test with only empty tensors
            values = [empty for i in range(3)]
            traced = torch.jit.trace(foo, values)
            x = warmup_and_run_forward(traced, *values)
            self.assertLastGraphAllFused()
            ref = foo(*values)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

    def test_cat_with_constant_dim(self):
        for device in self.devices:
            def foo(*args):
                v1 = torch.cat(args, dim=1)
                v2 = torch.cat([v1], dim=1)
                return v2 * v2

            empty = torch.tensor([], device=device, dtype=torch.float32)
            inputs = [empty] + [torch.randn(1, 64, device=device), torch.randn(1, 64, device=device)]
            traced = torch.jit.trace(foo, inputs)

            x = warmup_and_run_forward(traced, *inputs)
            self.assertLastGraphAllFused()
            ref = foo(*inputs)
            np.testing.assert_allclose(ref.cpu().numpy(), x.cpu().numpy())

    def test_scalar(self):
        @torch.jit.script
        def test_float(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor, a: float, b: float) -> torch.Tensor:
            return torch.add(torch.add(x, y, alpha=a), z, alpha=b)

        @torch.jit.script
        def test_int(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor, a: int, b: int) -> torch.Tensor:
            return torch.add(torch.add(x, y, alpha=a), z, alpha=b)

        for test in (test_float, test_int):
            for data_type in self.dtypes:
                x, y, z = (torch.rand(4, dtype=data_type) for i in range(3))
                a, b = 1, 2
                test(x, y, z, a, b)
                r = test(x, y, z, a, b)
                self.assertEqual(r, x + y * a + z * b)

    def test_loop(self):
        @torch.jit.script
        def test(x: torch.Tensor, y: torch.Tensor, z: int) -> torch.Tensor:
            b = y
            for i in range(0, z):
                a = x + y
                b = b + y
            return b

        x, y, z = (torch.zeros(32, 32), torch.ones(32, 32), 4)
        test(x, y, z)
        r = test(x, y, z)

    def test_slice(self):
        def easy(x, y):
            a = x[0:512:2]
            b = y[0:512:2]
            return a + b

        traced = torch.jit.trace(easy, (torch.ones(1024, 1024), torch.zeros(1024, 1024)))

        a = torch.ones(1024, 1024)
        x = traced(a, a)
        npr = a[0:512:2]
        npr = npr + npr
        np.testing.assert_allclose(npr.numpy(), x.numpy())

    def test_unsqueeze(self, N=256):
        def easy(x, y):
            a = torch.unsqueeze(x, 0)
            b = torch.unsqueeze(y, 0)
            return a + b

        traced = torch.jit.trace(easy, (torch.ones(N, N), torch.zeros(N, N)))

        a = torch.rand(N, N)
        x = traced(a, a)
        npr = np.expand_dims(a, 0)
        npr = npr + npr
        np.testing.assert_allclose(npr, x.numpy())

    def _test_softmax(self, device):
        def test_softmax(x, y):
            a = F.softmax(x, dim=0, dtype=torch.float32)
            b = F.softmax(y, dim=0, dtype=torch.float32)
            c = F.softmax(x, dim=1, dtype=torch.float32)
            d = F.softmax(y, dim=1, dtype=torch.float32)
            return a + b + c + d

        def test_softmax_neg_index(x, y):
            a = F.softmax(x, dim=-2, dtype=torch.float32)
            b = F.softmax(y, dim=-2, dtype=torch.float32)
            c = F.softmax(x, dim=-1, dtype=torch.float32)
            d = F.softmax(y, dim=-1, dtype=torch.float32)
            return a + b + c + d

        def test_log_softmax(x, y):
            a = F.log_softmax(x, dim=0, dtype=torch.float32)
            b = F.log_softmax(y, dim=0, dtype=torch.float32)
            c = F.log_softmax(x, dim=1, dtype=torch.float32)
            d = F.log_softmax(y, dim=1, dtype=torch.float32)
            return a + b + c + d

        for test in (test_softmax, test_log_softmax, test_softmax_neg_index):
            for data_type in self.dtypes:
                old = torch._C._jit_set_texpr_reductions_enabled(True)
                traced_input = torch.randn(2, 3, dtype=data_type, device=device)
                traced = torch.jit.trace(test, (traced_input, traced_input))
                inp = torch.randn(2, 3, dtype=data_type, device=device)
                res = traced(inp, inp)
                # Use eager mode as reference.
                ref = test(inp, inp)
                np.testing.assert_allclose(ref, res.cpu().numpy(), rtol=1e-06, atol=1e-06)
                torch._C._jit_set_texpr_reductions_enabled(old)

    def test_softmax_cpu(self):
        self._test_softmax('cpu')

    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
    @unittest.skip("global allocs are not supported yet.")
    def test_softmax_cuda(self):
        self._test_softmax('cuda')

    def test_half_gelu(self):
        devices = ["cuda"] if torch.cuda.is_available() else []

        @torch.jit.script
        def bias_gelu(bias, y):
            x = bias + y
            return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

        for device in devices:
            a = torch.rand(1024, dtype=torch.half, device=device)
            b = torch.rand(1024, dtype=torch.half, device=device)
            traced = torch.jit.trace(bias_gelu, (a, b))
            x = warmup_and_run_forward(traced, a, b)
            self.assertLastGraphAllFused()

    def test_half_bn_relu(self):
        devices = ["cuda"] if torch.cuda.is_available() else []

        def foo(a, b, c):
            y = torch.nn.functional.batch_norm(a, b, c)
            z = y.relu()
            return z

        for device in devices:
            a = torch.rand(16, 16, dtype=torch.half, device=device)
            b = torch.rand(16, dtype=torch.half, device=device)
            c = torch.rand(16, dtype=torch.half, device=device)
            traced = torch.jit.trace(foo, (a, b, c))
            print(traced.graph)
            x = warmup_and_run_forward(traced, a, b, c)
            self.assertLastGraphAllFused()

    def test_exp_pow(self):
        @torch.jit.script
        def do_exp(x, y, z):
            return ((x * y) * 2) * torch.pow(z, 2)

        for device in self.devices:
            x = torch.rand(10, dtype=torch.double, device=device)
            y = torch.rand(10, dtype=torch.double, device=device)
            z = torch.rand(10, dtype=torch.double, device=device)
            traced = torch.jit.trace(do_exp, (x, y, z))
            x = warmup_and_run_forward(traced, x, y, z)
            self.assertLastGraphAllFused()

    def test_sin_pow(self):
        def test(x):
            return torch.sin(torch.pow(x, 0))

        for data_type, shape in itertools.product(self.dtypes, [[3], [5], [10]]):
            x = torch.rand(shape, dtype=data_type)
            scripted = torch.jit.script(test)
            out = warmup_and_run_forward(scripted, x)
            self.assertLastGraphAllFused()
            self.assertEqual(out, test(x))

    def test_transpose(self):
        @torch.jit.script
        def test(x, y, z):
            return x.transpose(0, 1) + y + z
        x = torch.rand(4, 5, 2, 3)
        y = torch.rand(5, 4, 2, 3)
        z = torch.rand(5, 4, 2, 3)
        ref = test(x, y, z)
        res = test(x, y, z)
        np.testing.assert_allclose(ref.numpy(), res.numpy())

    def test_sliced_stride(self):
        @torch.jit.script
        def test(x, y, z):
            return x + y + z
        x = torch.rand(16, 4, 2, 3)[::2]
        y = torch.rand(8, 4, 2, 3)
        z = torch.rand(8, 4, 2, 3)
        ref = test(x, y, z)
        res = test(x, y, z)
        np.testing.assert_allclose(ref.numpy(), res.numpy())

    @unittest.skip("dynamic shapes are not quite there yet")
    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
    def test_dynamic_shape(self):
        with num_profiled_runs(2):
            @torch.jit.script
            def test(x, y, z):
                return x * y * z
            x, y, z = (torch.rand(4, 8).cuda() for _ in range(3))
            ref = test(x, y, z)
            _ = test(*[torch.rand(6, 8).cuda() for _ in range(3)])
            res = test(x, y, z)
            np.testing.assert_allclose(ref.cpu().numpy(), res.cpu().numpy())

            # A wild broadcast appears.
            x = torch.rand(4, 8).cuda()
            y = torch.rand(1, 8).cuda()
            z = torch.rand(4, 1).cuda()
            res = test(x, y, z)
            xn, yn, zn = (t.cpu().numpy() for t in (x, y, z))
            np.testing.assert_allclose(res.cpu().numpy(), xn * yn * zn)

            # Mismatched shapes shouldn't reach codegen.
            x = torch.rand(4, 8).cuda()
            y = torch.rand(4, 8).cuda()
            z = torch.rand(5, 8).cuda()
            try:
                res = test(x, y, z)
            except RuntimeError as e:
                assert "The size of tensor a (4) must match" in e.args[0]

            # Changing a static dimension fails guards.
            # x, y, z = [torch.rand(4, 7).cuda() for _ in range(3)]
            # xn, yn, zn = [t.cpu().numpy() for t in (x, y, z)]
            # res = test(x, y, z)
            # print(test.graph_for(x, y, z))
            # np.testing.assert_allclose(res.cpu().numpy(), xn * yn * zn)

    @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA")
    def test_guard_fails(self):
        @torch.jit.script
        def test(x, y, z):
            return x * y * z
        r1 = test(*[torch.rand(4).cuda() for _ in range(3)])
        r2 = test(*[torch.rand(4).cuda() for _ in range(3)])
        r3 = test(*[torch.rand(4).cuda() for _ in range(3)])
        r4 = test(*[torch.rand(7).cuda() for _ in range(3)])

    def test_bitwise_ops(self):
        def run_and(x, y):
            return x & (x & y)

        def run_or(x, y):
            return x & (x | y)

        def run_xor(x, y):
            return x ^ (x ^ y)

        def run_lshift(x, y):
            return x & (x << y)

        def run_rshift(x, y):
            return x & (x >> y)

        fns = {run_and, run_or, run_xor, run_lshift, run_rshift}

        for device in self.devices:
            for fn in fns:
                a = torch.ones(128, dtype=torch.int32, device=device)
                b = torch.zeros(128, dtype=torch.int32, device=device)
                inp = torch.ones(128, dtype=torch.int32, device=device)
                traced = torch.jit.trace(fn, (inp, inp))
                x = warmup_and_run_forward(traced, a, b)
                self.assertLastGraphAllFused()
                y = fn(a, b)
                np.testing.assert_allclose(x.cpu().numpy(), y.cpu().numpy())

    def test_where(self):
        def run_where(x, y):
            return torch.where(torch.gt(x, y), x, y)

        for data_type in self.dtypes:
            a = torch.rand(1024, dtype=data_type)
            b = torch.rand(1024, dtype=data_type)
            zeros = torch.zeros(1024, dtype=data_type)
            traced = torch.jit.trace(run_where, (zeros, zeros))
            x = warmup_and_run_forward(traced, a, b)
            self.assertLastGraphAllFused()
            y = run_where(a, b)
            np.testing.assert_allclose(x.float().numpy(), y.float().numpy())

    def test_multi_rand(self):
        for device in self.devices:
            def test(x):
                y = torch.rand_like(x)
                return (x + y) - (y - x)

            _atol = 2e-3
            _rtol = 1e-5
            for data_type in self.dtypes:
                if data_type is torch.bfloat16:
                    _atol = 2e-2
                a = torch.rand(4, dtype=data_type, device=device)
                scripted = torch.jit.script(test)
                out = warmup_and_run_forward(scripted, a)
                self.assertLastGraphAllFused()
                assert torch.allclose(out, 2 * a, atol=_atol, rtol=_rtol)

    def test_mask(self):
        def test(x):
            return x.unsqueeze(1) == 0

        for d in self.devices:
            for data_type in self.dtypes:
                x = torch.rand(4, dtype=data_type, device=d) > 0.5
                scripted = torch.jit.script(test)
                out = warmup_and_run_forward(scripted, x)
                self.assertLastGraphAllFused()
                assert torch.equal(out, test(x))

    def test_simple_add(self):
        val = torch._C._jit_get_te_generate_block_code()
        torch._C._jit_set_te_generate_block_code(True)
        fall_bk = torch._C._jit_texpr_fallback_allowed()
        torch._C._jit_texpr_set_fallback_allowed(True)

        def simple(a, b):
            return torch.add(a, b)

        a = torch.ones(256, 256)
        b = torch.ones(256, 256)
        traced = torch.jit.trace(simple,
                                 (torch.ones(256, 256), torch.ones(256, 256)))
        f = traced(a, b)
        f_test = np.full((256, 256), 2, dtype=float)
        np.testing.assert_allclose(f.numpy(), f_test)
        torch._C._jit_set_te_generate_block_code(val)
        torch._C._jit_texpr_set_fallback_allowed(fall_bk)

    def test_strided_output_preserved(self):
        def foo(a, b):
            return a + b - a

        # smaller, easier to debug example
        x = torch.arange(6)
        x = torch.as_strided(x, (2, 3), (1, 2))
        total = 0
        for i in range(2):
            for j in range(3):
                x[i, j] = total
                total += 1
        foo_script = torch.jit.script(foo)
        foo_script(x, x)
        foo_script(x, x)
        out_s = foo_script(x, x)
        out_eager = foo(x, x)
        self.assertEqual(out_s, out_eager)
        self.assertEqual(out_s.stride(), out_eager.stride())
        self.assertLastGraphAllFused()

        # more dims
        N, C, H, W, = 2, 3, 4, 5
        x = torch.rand(N, C, H, W).to(memory_format=torch.channels_last)
        foo_script = torch.jit.script(foo)
        foo_script(x, x)
        foo_script(x, x)
        out_s = foo_script(x, x)
        out_eager = foo(x, x)
        self.assertEqual(out_s, out_eager)
        self.assertEqual(out_s.stride(), out_eager.stride())
        self.assertLastGraphAllFused()

    def test_alias_analysis_module(self):
        class AliasModule(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                z = z + self.a
                self.b.add_(y)
                w = z + self.a
                z = w + x
                return z
        x = torch.randn(128, 128)

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am

        am = getModule(False)
        am_s = getModule(True)
        ref = am(x, x, x)
        test = am_s(x, x, x)
        torch.testing.assert_close(ref, test)

        # Now do the aliasing
        am.a = am.b
        ref = am(x, x, x)

        am_s.a = am_s.b
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)

    def test_alias_analysis_inputs(self):
        class AliasModule(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.a
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)

    def test_alias_analysis_input_and_module(self):
        class AliasModule(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                torch.manual_seed(1337)
                self.a = torch.randn(128, 128)
                self.b = torch.randn(128, 128)
                self.c = torch.randn(128, 128)

            def forward(self, x, y, z):
                x.add_(y)
                w = z + self.b
                z = w + x
                return z

        def getModule(script):
            am = AliasModule()
            if script:
                return torch.jit.script(am)
            return am
        am = getModule(False)
        am_s = getModule(True)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am.b = x
        ref = am(x, x, x)

        torch.manual_seed(1337)
        x = torch.randn(128, 128)
        am_s.b = x
        test = am_s(x, x, x)

        torch.testing.assert_close(ref, test)

    def test_multiple_outputs(self):
        for device in self.devices:
            # A bug reported internally similar to the one reported in #48533
            def foo(a, b, c):
                t_next = c + 1
                t5 = t_next * b
                t6 = torch.unsqueeze(t_next, 1)
                t7 = a * t6
                return (t7, t5, t_next)

            for data_type in self.dtypes:
                a = torch.rand(20, 20, dtype=data_type, device=device)
                b = torch.rand(20 * 29, dtype=data_type, device=device).as_strided([20], [29])
                c = torch.ones(20, dtype=torch.int64, device=device)
                traced = torch.jit.trace(foo, (a, b, c))
                ref = foo(a, b, c)
                exp = traced(a, b, c)
                exp = traced(a, b, c)
                self.assertEqual(ref, exp)

    def test_propagated_mem_layout(self):
        def foo(a, b, c):
            t_next = c + 1
            t5 = t_next * b
            t7 = a * t5
            return t7

        def foo_multi_outputs(a, b, c):
            t_next = c + 1
            t5 = b * t_next
            t7 = a * t5
            return (t7, t5, t_next)

        def foo_multi_outputs_i_nhwc_o_nchw(a, b, c):
            t_next = c + 1
            t5 = b * t_next
            t7 = a * t5
            t8 = t7.to(memory_format=torch.contiguous_format)
            return (t8, t7, t5, t_next)

        def run_foo_case(foo, a, b, c):
            traced_contiguous = torch.jit.trace(foo, (a, b, c))
            ref = foo(a, b, c)
            exp = traced_contiguous(a, b, c)
            exp = traced_contiguous(a, b, c)
            self.assertEqual(ref, exp)

        mem_layouts = list(itertools.product([torch.contiguous_format, torch.channels_last], repeat=3))
        shapes = [(2, 3, 4, 5), (2, 1, 1, 5), (1, 1, 1, 1)]
        permutes = [(0, 3, 2, 1), (0, 3, 1, 2)]
        funcs = [foo, foo_multi_outputs, foo_multi_outputs_i_nhwc_o_nchw]
        configs = itertools.product(funcs, shapes, mem_layouts, permutes)
        for strategy in ["STATIC", "DYNAMIC"]:
            old_strategy = torch.jit.set_fusion_strategy([(strategy, 10)])
            for _func, _shape, _mem_layouts, _permute in configs:
                a = torch.rand(_shape, dtype=torch.float32).to(memory_format=_mem_layouts[0])
                b = torch.rand(_shape, dtype=torch.float32).to(memory_format=_mem_layouts[1])
                c = torch.rand(_shape, dtype=torch.float32).to(memory_format=_mem_layouts[2])
                run_foo_case(_func, a, b, c)

                a = a.permute(dims=_permute)
                b = b.permute(dims=_permute)
                c = c.permute(dims=_permute)
                run_foo_case(_func, a, b, c)

            torch.jit.set_fusion_strategy(old_strategy)

if __name__ == '__main__':
    run_tests()