        torch.randn(2, 3, device="cuda"),
        torch.randn(2, 3, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        t1_b = fd.ops.broadcast_in_dim(t1, t0.shape(), [0, 1])
        t2 = fd.ops.add(t0, t1_b)

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = refs.add(
        inputs[0], prims.broadcast_in_dim(inputs[1], inputs[0].size(), [0, 1])
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_tensor_shape_with_output_bcast(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True])

        t1 = fd.ops.sum(t0, dims=[2])
        t1_b = fd.ops.broadcast_in_dim(t1, t0.shape(), [0, 1])

        fd.add_output(t1_b)

    inputs_1 = [
        torch.randn(2, 3, 4, device="cuda"),
    ]

    inputs_2 = [
        torch.randn(4, 5, 32, device="cuda"),
    ]

    inputs = inputs_1
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = prims.broadcast_in_dim(
        torch.sum(inputs[0], dim=-1), inputs[0].size(), [0, 1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    # Testing Dynamic usage of same Fusion
    inputs = inputs_2
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, new_fusion_expected=False
    )
    eager_out = prims.broadcast_in_dim(
        torch.sum(inputs[0], dim=-1), inputs[0].size(), [0, 1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_tensor_size_both_args_bcast(nvfuser_direct_test):
    inputs = [
        torch.randn(1, 3, device="cuda"),
        torch.randn(2, 1, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        t0_b = fd.ops.broadcast_in_dim(t0, [t1.size(0), t0.size(1)], [0, 1])
        t1_b = fd.ops.broadcast_in_dim(t1, [t1.size(0), t0.size(1)], [0, 1])
        t2 = fd.ops.add(t0_b, t1_b)

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = refs.add(


    ]

    inputs = inputs_1
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = prims.broadcast_in_dim(
        torch.sum(inputs[0], dim=-1), inputs[0].size(), [0, 1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    # Testing Dynamic usage of same Fusion
    inputs = inputs_2
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, new_fusion_expected=False
    )
    eager_out = prims.broadcast_in_dim(
        torch.sum(inputs[0], dim=-1), inputs[0].size(), [0, 1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_tensor_size_both_args_bcast(nvfuser_direct_test):
    inputs = [
        torch.randn(1, 3, device="cuda"),
        torch.randn(2, 1, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        t0_b = fd.ops.broadcast_in_dim(t0, [t1.size(0), t0.size(1)], [0, 1])
        t1_b = fd.ops.broadcast_in_dim(t1, [t1.size(0), t0.size(1)], [0, 1])
        t2 = fd.ops.add(t0_b, t1_b)

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = refs.add(
        prims.broadcast_in_dim(
            inputs[0], [inputs[1].size()[0], inputs[0].size()[1]], [0, 1]
        ),
        prims.broadcast_in_dim(
            inputs[1], [inputs[1].size()[0], inputs[0].size()[1]], [0, 1]
        ),
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_var_correction(nvfuser_direct_test):
    num_elem = 2
    inputs = [torch.randn(2, num_elem, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.var(t0, [-1], correction)
            fd.add_output(t1)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs


        t1 = fd.from_pytorch(inputs[1])

        t0_b = fd.ops.broadcast_in_dim(t0, [t1.size(0), t0.size(1)], [0, 1])
        t1_b = fd.ops.broadcast_in_dim(t1, [t1.size(0), t0.size(1)], [0, 1])
        t2 = fd.ops.add(t0_b, t1_b)

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = refs.add(
        prims.broadcast_in_dim(
            inputs[0], [inputs[1].size()[0], inputs[0].size()[1]], [0, 1]
        ),
        prims.broadcast_in_dim(
            inputs[1], [inputs[1].size()[0], inputs[0].size()[1]], [0, 1]
        ),
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_var_correction(nvfuser_direct_test):
    num_elem = 2
    inputs = [torch.randn(2, num_elem, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.var(t0, [-1], correction)
            fd.add_output(t1)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var(inputs[0], [-1], correction=correction)
        nvfuser_direct_test.assertEqual(fuser_result[0], torch_result)

def test_var_mean_correction(nvfuser_direct_test):
    num_elem = 2
    inputs = [torch.randn(2, num_elem, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [-1], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(


    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.var(t0, [-1], correction)
            fd.add_output(t1)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var(inputs[0], [-1], correction=correction)
        nvfuser_direct_test.assertEqual(fuser_result[0], torch_result)

def test_var_mean_correction(nvfuser_direct_test):
    num_elem = 2
    inputs = [torch.randn(2, num_elem, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [-1], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [-1], correction=correction)
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_zero_size_dim(nvfuser_direct_test):
    inputs = [
        torch.ones(0, 0, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(
            shape=[0, 0], contiguity=[True, True], dtype=DataType.Float
        )
        t1 = fd.ops.relu(t0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.relu(inputs[0])
    nvfuser_direct_test.assertEqual(eager_out.numel(), nvf_out[0].numel())

def test_allocation_domain_concretization(nvfuser_direct_test):


    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [-1], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    # correction must be less than the reduction factor, which is the input
    # numel divided by output numel.
    for correction in range(num_elem):
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [-1], correction=correction)
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_zero_size_dim(nvfuser_direct_test):
    inputs = [
        torch.ones(0, 0, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(
            shape=[0, 0], contiguity=[True, True], dtype=DataType.Float
        )
        t1 = fd.ops.relu(t0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.relu(inputs[0])
    nvfuser_direct_test.assertEqual(eager_out.numel(), nvf_out[0].numel())

def test_allocation_domain_concretization(nvfuser_direct_test):
    inputs = [
        # we need an empty tensor here so we'll trigger `concretizeEmptyExtents`
        torch.randn((0,), dtype=torch.float64, device="cuda:0").as_strided(
            (1, 0, 1, 1), (0, 1, 1, 1)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[1, -1, 1, 1],
            contiguity=[True, None, None, None],
            dtype=DataType.Double,
            is_cpu=False,
            stride_order=[0, 3, 2, 1],
        )
        S1 = fd.define_scalar(2.0, dtype=DataType.Double)
        T2 = fd.ops.mul(T1, S1)


        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_zero_size_dim(nvfuser_direct_test):
    inputs = [
        torch.ones(0, 0, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(
            shape=[0, 0], contiguity=[True, True], dtype=DataType.Float
        )
        t1 = fd.ops.relu(t0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.relu(inputs[0])
    nvfuser_direct_test.assertEqual(eager_out.numel(), nvf_out[0].numel())

def test_allocation_domain_concretization(nvfuser_direct_test):
    inputs = [
        # we need an empty tensor here so we'll trigger `concretizeEmptyExtents`
        torch.randn((0,), dtype=torch.float64, device="cuda:0").as_strided(
            (1, 0, 1, 1), (0, 1, 1, 1)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[1, -1, 1, 1],
            contiguity=[True, None, None, None],
            dtype=DataType.Double,
            is_cpu=False,
            stride_order=[0, 3, 2, 1],
        )
        S1 = fd.define_scalar(2.0, dtype=DataType.Double)
        T2 = fd.ops.mul(T1, S1)
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    torch_ref = inputs[0] * 2.0
    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_allocation_domain_index_select(nvfuser_direct_test):
    inputs = [
        torch.randn((252,), dtype=torch.float32, device="cuda:0").as_strided(
            (9, 28), (1, 9)
        ),
        torch.randint(0, 28, (4,), dtype=torch.int64, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0, 1],
        )
        T2 = fd.define_tensor(
            shape=[-1], contiguity=[True], dtype=DataType.Int, is_cpu=False


        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[1, -1, 1, 1],
            contiguity=[True, None, None, None],
            dtype=DataType.Double,
            is_cpu=False,
            stride_order=[0, 3, 2, 1],
        )
        S1 = fd.define_scalar(2.0, dtype=DataType.Double)
        T2 = fd.ops.mul(T1, S1)
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    torch_ref = inputs[0] * 2.0
    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_allocation_domain_index_select(nvfuser_direct_test):
    inputs = [
        torch.randn((252,), dtype=torch.float32, device="cuda:0").as_strided(
            (9, 28), (1, 9)
        ),
        torch.randint(0, 28, (4,), dtype=torch.int64, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0, 1],
        )
        T2 = fd.define_tensor(
            shape=[-1], contiguity=[True], dtype=DataType.Int, is_cpu=False
        )
        T3 = fd.ops.index_select(T1, T2, dim=1)
        fd.add_output(T3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    torch_ref = torch.index_select(inputs[0], 1, inputs[1])
    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_expand_to_zero(nvfuser_direct_test):
    inputs = [
        # This is an actually empty tensor
        torch.zeros((1, 0), dtype=torch.float32, device="cuda:0"),
        # This one is not actually empty, but should appear to be empty due to expand
        torch.zeros((1, 1), dtype=torch.float32, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.broadcast_in_dim(T0, shape=[0, 0], broadcast_dims=[0, 1])
        T3 = fd.ops.broadcast_in_dim(T1, shape=[0, 0], broadcast_dims=[0, 1])
        fd.add_output(T2)
        fd.add_output(T3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)


    def fusion_func(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0, 1],
        )
        T2 = fd.define_tensor(
            shape=[-1], contiguity=[True], dtype=DataType.Int, is_cpu=False
        )
        T3 = fd.ops.index_select(T1, T2, dim=1)
        fd.add_output(T3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    torch_ref = torch.index_select(inputs[0], 1, inputs[1])
    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_expand_to_zero(nvfuser_direct_test):
    inputs = [
        # This is an actually empty tensor
        torch.zeros((1, 0), dtype=torch.float32, device="cuda:0"),
        # This one is not actually empty, but should appear to be empty due to expand
        torch.zeros((1, 1), dtype=torch.float32, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.broadcast_in_dim(T0, shape=[0, 0], broadcast_dims=[0, 1])
        T3 = fd.ops.broadcast_in_dim(T1, shape=[0, 0], broadcast_dims=[0, 1])
        fd.add_output(T2)
        fd.add_output(T3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    nvfuser_direct_test.assertEqual(nvf_out[0].shape, (0, 0))
    nvfuser_direct_test.assertEqual(nvf_out[1].shape, (0, 0))

def test_expanded_bcast_tensor(nvfuser_direct_test):
    inputs = [
        torch.tensor(1.5, device="cuda"),
        torch.randn(5, 5, 5, device="cuda"),
        torch.randint(0, 1, (5, 5), device="cuda").bool().unsqueeze(-1).expand(5, 5, 5),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.add(T0, T1)
        T4 = fd.ops.add(T2, T3)
        fd.add_output(T4)

    eager_out = inputs[0] + inputs[1] + inputs[2]

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)


        # This is an actually empty tensor
        torch.zeros((1, 0), dtype=torch.float32, device="cuda:0"),
        # This one is not actually empty, but should appear to be empty due to expand
        torch.zeros((1, 1), dtype=torch.float32, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.broadcast_in_dim(T0, shape=[0, 0], broadcast_dims=[0, 1])
        T3 = fd.ops.broadcast_in_dim(T1, shape=[0, 0], broadcast_dims=[0, 1])
        fd.add_output(T2)
        fd.add_output(T3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    nvfuser_direct_test.assertEqual(nvf_out[0].shape, (0, 0))
    nvfuser_direct_test.assertEqual(nvf_out[1].shape, (0, 0))

def test_expanded_bcast_tensor(nvfuser_direct_test):
    inputs = [
        torch.tensor(1.5, device="cuda"),
        torch.randn(5, 5, 5, device="cuda"),
        torch.randint(0, 1, (5, 5), device="cuda").bool().unsqueeze(-1).expand(5, 5, 5),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.add(T0, T1)
        T4 = fd.ops.add(T2, T3)
        fd.add_output(T4)

    eager_out = inputs[0] + inputs[1] + inputs[2]

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_inplace_update_on_non_contiguous_inputs(nvfuser_direct_test):
    inputs = [
        torch.randn(5, dtype=torch.float32, device="cuda:0").as_strided((2, 2), (1, 3)),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[2, 2],
            contiguity=[False, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0, 1],
        )
        S1 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T2 = fd.ops.gt(T0, S1)
        S3 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T4 = fd.ops.where(T2, T0, S3)
        T5 = fd.ops.cast(T4, dtype=DataType.Float)


    inputs = [
        torch.tensor(1.5, device="cuda"),
        torch.randn(5, 5, 5, device="cuda"),
        torch.randint(0, 1, (5, 5), device="cuda").bool().unsqueeze(-1).expand(5, 5, 5),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.add(T0, T1)
        T4 = fd.ops.add(T2, T3)
        fd.add_output(T4)

    eager_out = inputs[0] + inputs[1] + inputs[2]

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_inplace_update_on_non_contiguous_inputs(nvfuser_direct_test):
    inputs = [
        torch.randn(5, dtype=torch.float32, device="cuda:0").as_strided((2, 2), (1, 3)),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[2, 2],
            contiguity=[False, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0, 1],
        )
        S1 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T2 = fd.ops.gt(T0, S1)
        S3 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T4 = fd.ops.where(T2, T0, S3)
        T5 = fd.ops.cast(T4, dtype=DataType.Float)
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T0)
        fd.add_output(T6)

    ref_inp = inputs[0].clone()
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func,
        inputs,
    )

    assert len(nvf_out) == 1
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0])
    nvfuser_direct_test.assertEqual(nvf_out[0], ref_inp.relu())

def test_pad_expanded_empty(nvfuser_direct_test):
    inputs = [
        torch.randn((0,), dtype=torch.float64, device="cuda:0").as_strided(
            (2, 0, 3), (0, 0, 0)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        S1 = fd.define_scalar(-3.70753, dtype=DataType.Double)
        T2 = fd.ops.pad(T0, [0, 0, 1, 1, 1, 0], S1)
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    torch_ref = torch.nn.functional.pad(
        inputs[0], (0, 0, 1, 1, 1, 0), "constant", -3.70753
    )


        S1 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T2 = fd.ops.gt(T0, S1)
        S3 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T4 = fd.ops.where(T2, T0, S3)
        T5 = fd.ops.cast(T4, dtype=DataType.Float)
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T0)
        fd.add_output(T6)

    ref_inp = inputs[0].clone()
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func,
        inputs,
    )

    assert len(nvf_out) == 1
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0])
    nvfuser_direct_test.assertEqual(nvf_out[0], ref_inp.relu())

def test_pad_expanded_empty(nvfuser_direct_test):
    inputs = [
        torch.randn((0,), dtype=torch.float64, device="cuda:0").as_strided(
            (2, 0, 3), (0, 0, 0)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        S1 = fd.define_scalar(-3.70753, dtype=DataType.Double)
        T2 = fd.ops.pad(T0, [0, 0, 1, 1, 1, 0], S1)
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    torch_ref = torch.nn.functional.pad(
        inputs[0], (0, 0, 1, 1, 1, 0), "constant", -3.70753
    )

    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_pad_prior_cat(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(3, 3, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        # pad tensors t0 and t1, so their first dimension are size 10.
        t0_pad = fd.ops.pad(t0, [0, 0, 0, 8])
        t1_pad = fd.ops.pad(t1, [0, 0, 0, 7])

        t3 = fd.ops.cat([t0_pad, t1_pad], 1)
        fd.add_output(t3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)


        torch.randn((0,), dtype=torch.float64, device="cuda:0").as_strided(
            (2, 0, 3), (0, 0, 0)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        S1 = fd.define_scalar(-3.70753, dtype=DataType.Double)
        T2 = fd.ops.pad(T0, [0, 0, 1, 1, 1, 0], S1)
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    torch_ref = torch.nn.functional.pad(
        inputs[0], (0, 0, 1, 1, 1, 0), "constant", -3.70753
    )

    nvfuser_direct_test.assertEqual(nvf_out[0], torch_ref)

def test_pad_prior_cat(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(3, 3, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        # pad tensors t0 and t1, so their first dimension are size 10.
        t0_pad = fd.ops.pad(t0, [0, 0, 0, 8])
        t1_pad = fd.ops.pad(t1, [0, 0, 0, 7])

        t3 = fd.ops.cat([t0_pad, t1_pad], 1)
        fd.add_output(t3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    # pad tensors t0 and t1, so their first dimension are size 10.
    pad_input0 = torch.nn.functional.pad(inputs[0], [0, 0, 0, 8])
    pad_input1 = torch.nn.functional.pad(inputs[1], [0, 0, 0, 7])
    nvfuser_direct_test.assertEqual(
        torch.cat([pad_input0, pad_input1], dim=1), nvf_out[0]
    )

def test_replaced_sizes_pr2714(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.exp(T0)
        T3 = fd.ops.tanh(T1)


        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        # pad tensors t0 and t1, so their first dimension are size 10.
        t0_pad = fd.ops.pad(t0, [0, 0, 0, 8])
        t1_pad = fd.ops.pad(t1, [0, 0, 0, 7])

        t3 = fd.ops.cat([t0_pad, t1_pad], 1)
        fd.add_output(t3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    # pad tensors t0 and t1, so their first dimension are size 10.
    pad_input0 = torch.nn.functional.pad(inputs[0], [0, 0, 0, 8])
    pad_input1 = torch.nn.functional.pad(inputs[1], [0, 0, 0, 7])
    nvfuser_direct_test.assertEqual(
        torch.cat([pad_input0, pad_input1], dim=1), nvf_out[0]
    )

def test_replaced_sizes_pr2714(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.exp(T0)
        T3 = fd.ops.tanh(T1)
        S4 = fd.define_scalar(4, dtype=DataType.Int)
        T6 = fd.ops.reshape(T2, new_shape=[S4])
        S7 = fd.define_scalar(4, dtype=DataType.Int)
        T9 = fd.ops.reshape(T3, new_shape=[S7])
        T10 = fd.ops.add(T6, T9)
        T11 = fd.ops.reciprocal(T0)
        T12 = fd.ops.mul(T3, T11)
        S13 = fd.define_scalar(2.00000, dtype=DataType.Double)
        S14 = fd.ops.reciprocal(S13)
        T15 = fd.ops.mul(T10, S14)
        fd.add_output(T10)
        fd.add_output(T12)
        fd.add_output(T15)

    inputs = [
        torch.randn((4,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 2), (2, 1)
        ),
        torch.randn((4,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 2), (2, 1)
        ),
    ]

    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_reshape_squeeze_concretization(nvfuser_direct_test):
    inputs = [
        torch.randn((100,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 5, 10), (50, 10, 1)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T1 = fd.ops.slice(
            T0, start_indices=[0, 0, 0], end_indices=[1, 2, 4], strides=[1, 1, 1]
        )


        T12 = fd.ops.mul(T3, T11)
        S13 = fd.define_scalar(2.00000, dtype=DataType.Double)
        S14 = fd.ops.reciprocal(S13)
        T15 = fd.ops.mul(T10, S14)
        fd.add_output(T10)
        fd.add_output(T12)
        fd.add_output(T15)

    inputs = [
        torch.randn((4,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 2), (2, 1)
        ),
        torch.randn((4,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 2), (2, 1)
        ),
    ]

    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_reshape_squeeze_concretization(nvfuser_direct_test):
    inputs = [
        torch.randn((100,), dtype=torch.float32, device="cuda:0").as_strided(
            (2, 5, 10), (50, 10, 1)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T1 = fd.ops.slice(
            T0, start_indices=[0, 0, 0], end_indices=[1, 2, 4], strides=[1, 1, 1]
        )
        S2 = fd.define_scalar(1, dtype=DataType.Int)
        S3 = fd.define_scalar(8, dtype=DataType.Int)
        T6 = fd.ops.reshape(T1, new_shape=[S2, S3])
        T7 = fd.ops.reshape(T6, new_shape=[S3])
        fd.add_output(T7)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_sum_sliced_reshape_to_broadcast(nvfuser_direct_test):
    inputs = [torch.randn((24, 128, 25, 32), dtype=torch.float32, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T18 = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=[True, True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )
        S91 = fd.define_scalar(12, dtype=DataType.Int)
        S92 = fd.define_scalar(128, dtype=DataType.Int)
        S93 = fd.define_scalar(25, dtype=DataType.Int)
        S94 = fd.define_scalar(32, dtype=DataType.Int)
        S95 = fd.define_scalar(2, dtype=DataType.Int)
        T97 = fd.ops.reshape(T18, new_shape=[S91, S92, S93, S94, S95])
        T98 = fd.ops.slice(
            T97,


    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T1 = fd.ops.slice(
            T0, start_indices=[0, 0, 0], end_indices=[1, 2, 4], strides=[1, 1, 1]
        )
        S2 = fd.define_scalar(1, dtype=DataType.Int)
        S3 = fd.define_scalar(8, dtype=DataType.Int)
        T6 = fd.ops.reshape(T1, new_shape=[S2, S3])
        T7 = fd.ops.reshape(T6, new_shape=[S3])
        fd.add_output(T7)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_sum_sliced_reshape_to_broadcast(nvfuser_direct_test):
    inputs = [torch.randn((24, 128, 25, 32), dtype=torch.float32, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T18 = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=[True, True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )
        S91 = fd.define_scalar(12, dtype=DataType.Int)
        S92 = fd.define_scalar(128, dtype=DataType.Int)
        S93 = fd.define_scalar(25, dtype=DataType.Int)
        S94 = fd.define_scalar(32, dtype=DataType.Int)
        S95 = fd.define_scalar(2, dtype=DataType.Int)
        T97 = fd.ops.reshape(T18, new_shape=[S91, S92, S93, S94, S95])
        T98 = fd.ops.slice(
            T97,
            start_indices=[0, 0, 0, 0, 0],
            end_indices=[12, 128, 25, 32, 1],
            strides=[1, 1, 1, 1, 1],
        )
        T89 = fd.ops.sum(T98, dims=[4], keepdim=False, dtype=DataType.Null)
        fd.add_output(T89)

    # Check that define_scalar is not used in fd_str
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1, -1], contiguity=[True, True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.ops.reshape(tv0, new_shape=[12, 128, 25, 32, 2])
    tv2 = fd.ops.slice(tv1, start_indices=[0, 0, 0, 0, 0], end_indices=[12, 128, 25, 32, 1], strides=[1, 1, 1, 1, 1], manual_normalization=True)
    tv3 = fd.ops.squeeze(tv2, dims=[4])
    fd.add_output(tv3)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )

# See https://github.com/NVIDIA/Fuser/issues/3833
def test_bcast_squeeze_replace_aliased_output(nvfuser_direct_test):
    inputs = [
        torch.testing.make_tensor((1, 1, 576), dtype=torch.bfloat16, device="cuda:0"),
        torch.testing.make_tensor((1, 576), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, 1, 576],
            contiguity=[None, None, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T1 = fd.define_tensor(
            shape=[1, 576],
            contiguity=[None, True],


            end_indices=[12, 128, 25, 32, 1],
            strides=[1, 1, 1, 1, 1],
        )
        T89 = fd.ops.sum(T98, dims=[4], keepdim=False, dtype=DataType.Null)
        fd.add_output(T89)

    # Check that define_scalar is not used in fd_str
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1, -1], contiguity=[True, True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.ops.reshape(tv0, new_shape=[12, 128, 25, 32, 2])
    tv2 = fd.ops.slice(tv1, start_indices=[0, 0, 0, 0, 0], end_indices=[12, 128, 25, 32, 1], strides=[1, 1, 1, 1, 1], manual_normalization=True)
    tv3 = fd.ops.squeeze(tv2, dims=[4])
    fd.add_output(tv3)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )

# See https://github.com/NVIDIA/Fuser/issues/3833
def test_bcast_squeeze_replace_aliased_output(nvfuser_direct_test):
    inputs = [
        torch.testing.make_tensor((1, 1, 576), dtype=torch.bfloat16, device="cuda:0"),
        torch.testing.make_tensor((1, 576), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, 1, 576],
            contiguity=[None, None, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T1 = fd.define_tensor(
            shape=[1, 576],
            contiguity=[None, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T5 = fd.ops.reshape(T0, new_shape=[1, 576])
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T1)
        fd.add_output(T5)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    assert len(nvf_out) == 1
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0].squeeze(1))

def test_broadcast_and_stride_order(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 3, 4, dtype=torch.float32, device="cuda:0"),
    ]

    # Direct bindings does not support `add_output` with stride_order argument.
    # Instead, we use `stride_order` operation to set the stride order before
    # adding the output.
    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.ops.broadcast(T0, is_broadcast_dim=[False, True, False, False])
        T2 = fd.ops.stride_order(T1, stride_order=[0, 1, 2, 3])
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0].unsqueeze(1))
    nvfuser_direct_test.assertEqual(nvf_out[0].stride(), (1, 2, 2, 6))


            stride_order=[2, 1, 0],
        )
        T1 = fd.define_tensor(
            shape=[1, 576],
            contiguity=[None, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T5 = fd.ops.reshape(T0, new_shape=[1, 576])
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T1)
        fd.add_output(T5)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    assert len(nvf_out) == 1
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0].squeeze(1))

def test_broadcast_and_stride_order(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 3, 4, dtype=torch.float32, device="cuda:0"),
    ]

    # Direct bindings does not support `add_output` with stride_order argument.
    # Instead, we use `stride_order` operation to set the stride order before
    # adding the output.
    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.ops.broadcast(T0, is_broadcast_dim=[False, True, False, False])
        T2 = fd.ops.stride_order(T1, stride_order=[0, 1, 2, 3])
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0].unsqueeze(1))
    nvfuser_direct_test.assertEqual(nvf_out[0].stride(), (1, 2, 2, 6))

def test_right_shift_logical(nvfuser_direct_test):
    dtypes = [torch.int32, torch.int64]
    input = torch.tensor(
        [
            -1,
            -2147483648,
            1073741824,
            -64463884,
            -65968277,
            4042311,
            -98914167,
            5526216,
        ],
        device="cuda",
    )

    # expected_outputs given by jax.lax.shift_right_logical(inputs, 3)
    expected_outputs = [


def test_broadcast_and_stride_order(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 3, 4, dtype=torch.float32, device="cuda:0"),
    ]

    # Direct bindings does not support `add_output` with stride_order argument.
    # Instead, we use `stride_order` operation to set the stride order before
    # adding the output.
    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.ops.broadcast(T0, is_broadcast_dim=[False, True, False, False])
        T2 = fd.ops.stride_order(T1, stride_order=[0, 1, 2, 3])
        fd.add_output(T2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0].unsqueeze(1))
    nvfuser_direct_test.assertEqual(nvf_out[0].stride(), (1, 2, 2, 6))

def test_right_shift_logical(nvfuser_direct_test):
    dtypes = [torch.int32, torch.int64]
    input = torch.tensor(
        [
            -1,
            -2147483648,
            1073741824,
            -64463884,
            -65968277,
            4042311,
            -98914167,
            5526216,
        ],
        device="cuda",
    )

    # expected_outputs given by jax.lax.shift_right_logical(inputs, 3)
    expected_outputs = [
        torch.tensor(
            [
                536870911,
                268435456,
                134217728,
                528812926,
                528624877,
                505288,
                524506641,
                690777,
            ],
            dtype=torch.int32,
            device="cuda",
        ),
        torch.tensor(
            [
                2305843009213693951,
                2305843008945258496,
                134217728,
                2305843009205635966,
                2305843009205447917,
                505288,
                2305843009201329681,
                690777,
            ],
            dtype=torch.int64,
            device="cuda",
        ),
    ]

    for idx, dtype in enumerate(dtypes):
        current_input = input.to(dtype)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(3)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [current_input])
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_outputs[idx])

def test_right_shift_logical_sizeof_dtype(nvfuser_direct_test):
    dtypes = [torch.int32, torch.int64]
    input = torch.tensor(
        [
            -1,
            -2147483648,
            1073741824,
            -64463884,
            -65968277,
            4042311,
            -98914167,
            5526216,
        ],
        device="cuda",
    )

    for idx, dtype in enumerate(dtypes):
        current_input = input.to(dtype)


                690777,
            ],
            dtype=torch.int64,
            device="cuda",
        ),
    ]

    for idx, dtype in enumerate(dtypes):
        current_input = input.to(dtype)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(3)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [current_input])
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_outputs[idx])

def test_right_shift_logical_sizeof_dtype(nvfuser_direct_test):
    dtypes = [torch.int32, torch.int64]
    input = torch.tensor(
        [
            -1,
            -2147483648,
            1073741824,
            -64463884,
            -65968277,
            4042311,
            -98914167,
            5526216,
        ],
        device="cuda",
    )

    for idx, dtype in enumerate(dtypes):
        current_input = input.to(dtype)
        num_bits = 32 if (dtype == torch.int32) else 64

        # expected_outputs given by jax.lax.shift_right_logical(inputs, sizeof(dtype))
        # >>> jax.lax.shift_right_logical(input.to('cpu').numpy(), 32)
        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        # >>> jax.lax.shift_right_logical(input.to('cpu').numpy(), 64)
        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        expected_output = torch.zeros_like(current_input)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(None, dtype=DataType.Int)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_func, [current_input, num_bits]
        )
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_output)

def test_all_dim_var_mean(nvfuser_direct_test):
    inputs = [torch.randn(2, 2, 2, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [0, 1, 2], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    list_of_test_cases = [0, 1]
    for correction in list_of_test_cases:
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )


        # expected_outputs given by jax.lax.shift_right_logical(inputs, sizeof(dtype))
        # >>> jax.lax.shift_right_logical(input.to('cpu').numpy(), 32)
        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        # >>> jax.lax.shift_right_logical(input.to('cpu').numpy(), 64)
        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        expected_output = torch.zeros_like(current_input)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(None, dtype=DataType.Int)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_func, [current_input, num_bits]
        )
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_output)

def test_all_dim_var_mean(nvfuser_direct_test):
    inputs = [torch.randn(2, 2, 2, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [0, 1, 2], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    list_of_test_cases = [0, 1]
    for correction in list_of_test_cases:
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [0, 1, 2], bool(correction))
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_validate_precomputed_values(nvfuser_direct_test):
    # This test is from legacy test_nan.py
    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )

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

    outs, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func,
        [
            torch.randn((10,), dtype=torch.float32, device="cuda:0").as_strided(


    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [0, 1, 2], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    list_of_test_cases = [0, 1]
    for correction in list_of_test_cases:
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [0, 1, 2], bool(correction))
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_validate_precomputed_values(nvfuser_direct_test):
    # This test is from legacy test_nan.py
    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )

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

    outs, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func,
        [
            torch.randn((10,), dtype=torch.float32, device="cuda:0").as_strided(
                (2, 5), (5, 1)
            ),
            float("nan"),
        ],
    )
    # Comparing any number to NaN results in False.
    torch.testing.assert_close(outs[0].cpu(), torch.full((2, 5), False))

def test_cpu_add(nvfuser_direct_test):
    """
    Tests that CPU scalar tensor can be instantiated using fd.from_pytorch

    Migrated from `test_cpu_add` in legacy test_pointwise.py
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
        torch.randn(3, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.from_pytorch(inputs[1])
        t1 = fd.ops.add(t0, s0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)


            is_cpu=False,
        )

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

    outs, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func,
        [
            torch.randn((10,), dtype=torch.float32, device="cuda:0").as_strided(
                (2, 5), (5, 1)
            ),
            float("nan"),
        ],
    )
    # Comparing any number to NaN results in False.
    torch.testing.assert_close(outs[0].cpu(), torch.full((2, 5), False))

def test_cpu_add(nvfuser_direct_test):
    """
    Tests that CPU scalar tensor can be instantiated using fd.from_pytorch

    Migrated from `test_cpu_add` in legacy test_pointwise.py
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
        torch.randn(3, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.from_pytorch(inputs[1])
        t1 = fd.ops.add(t0, s0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0] + inputs[1])

def test_full_with_cpu_inputs(nvfuser_direct_test):
    """
    This example contains CPU scalar only fusion inputs.
    The `full` op does not take any fusion inputs but generates a
    CUDA tensor. This is a nvFuser supported fusion since the final
    output is a CUDA tensor.
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.full(shape=[2, 2], fill_value=s0, dtype=DataType.Float)
        t2 = fd.ops.mul(tv0, tv1)  # CPU scalar * CUDA tensor = CUDA tensor
        fd.add_output(t2)


    """
    Tests that CPU scalar tensor can be instantiated using fd.from_pytorch

    Migrated from `test_cpu_add` in legacy test_pointwise.py
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
        torch.randn(3, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.from_pytorch(inputs[1])
        t1 = fd.ops.add(t0, s0)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0] + inputs[1])

def test_full_with_cpu_inputs(nvfuser_direct_test):
    """
    This example contains CPU scalar only fusion inputs.
    The `full` op does not take any fusion inputs but generates a
    CUDA tensor. This is a nvFuser supported fusion since the final
    output is a CUDA tensor.
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.full(shape=[2, 2], fill_value=s0, dtype=DataType.Float)
        t2 = fd.ops.mul(tv0, tv1)  # CPU scalar * CUDA tensor = CUDA tensor
        fd.add_output(t2)

    # fd.validate expects inputs and outputs to be on CUDA device
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_input_forwarding_device(nvfuser_direct_test):
    """
    If fusion segment do not consist of any exprs, no kernel is
    launched and the output is on the correct device.
    """
    inputs = [torch.tensor(2.0, device="cpu", dtype=torch.float)]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        fd.add_output(tv0)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert out[0].is_cpu

def test_single_segment_multi_device():
    """
    Test single segment with CPU and CUDA outputs


    This example contains CPU scalar only fusion inputs.
    The `full` op does not take any fusion inputs but generates a
    CUDA tensor. This is a nvFuser supported fusion since the final
    output is a CUDA tensor.
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.full(shape=[2, 2], fill_value=s0, dtype=DataType.Float)
        t2 = fd.ops.mul(tv0, tv1)  # CPU scalar * CUDA tensor = CUDA tensor
        fd.add_output(t2)

    # fd.validate expects inputs and outputs to be on CUDA device
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_input_forwarding_device(nvfuser_direct_test):
    """
    If fusion segment do not consist of any exprs, no kernel is
    launched and the output is on the correct device.
    """
    inputs = [torch.tensor(2.0, device="cpu", dtype=torch.float)]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        fd.add_output(tv0)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert out[0].is_cpu

def test_single_segment_multi_device():
    """
    Test single segment with CPU and CUDA outputs
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
        torch.tensor(3.0, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.add(tv0, s0)
        tv2 = fd.from_pytorch(inputs[1])
        tv3 = fd.ops.add(tv1, tv2)
        fd.add_output(tv1)
        fd.add_output(tv2)


    # fd.validate expects inputs and outputs to be on CUDA device
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_input_forwarding_device(nvfuser_direct_test):
    """
    If fusion segment do not consist of any exprs, no kernel is
    launched and the output is on the correct device.
    """
    inputs = [torch.tensor(2.0, device="cpu", dtype=torch.float)]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        fd.add_output(tv0)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert out[0].is_cpu

def test_single_segment_multi_device():
    """
    Test single segment with CPU and CUDA outputs
    """
    inputs = [
        torch.tensor(2.0, device="cpu", dtype=torch.float),
        torch.tensor(3.0, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.add(tv0, s0)
        tv2 = fd.from_pytorch(inputs[1])
        tv3 = fd.ops.add(tv1, tv2)
        fd.add_output(tv1)
        fd.add_output(tv2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    with pytest.raises(
        RuntimeError, match="KernelExecutor does not support the Fusion provided."
    ):
        _ = fd.execute(inputs)

# Test that we properly handle packed type
@pytest.mark.skipif(
    is_pre_blackwell(), reason="Only supported on blackwell and newer devices."
)
def test_packed_fp4(nvfuser_direct_test):
    t0 = torch.rand(
        (
            1024,
            32,
        ),
        dtype=torch.float32,
        device="cuda:0",
    )
    # we'll just ignore the scaling factor, since we only want to test basic fp4 support
    t0_fp4, _ = pytorch_nvfp4_quantize(t0, 1.0)
    inputs = [t0_fp4]

    def fusion_func(fd: FusionDefinition):


        s0 = fd.define_scalar(3.0)
        tv1 = fd.ops.add(tv0, s0)
        tv2 = fd.from_pytorch(inputs[1])
        tv3 = fd.ops.add(tv1, tv2)
        fd.add_output(tv1)
        fd.add_output(tv2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    with pytest.raises(
        RuntimeError, match="KernelExecutor does not support the Fusion provided."
    ):
        _ = fd.execute(inputs)

# Test that we properly handle packed type
@pytest.mark.skipif(
    is_pre_blackwell(), reason="Only supported on blackwell and newer devices."
)
def test_packed_fp4(nvfuser_direct_test):
    t0 = torch.rand(
        (
            1024,
            32,
        ),
        dtype=torch.float32,
        device="cuda:0",
    )
    # we'll just ignore the scaling factor, since we only want to test basic fp4 support
    t0_fp4, _ = pytorch_nvfp4_quantize(t0, 1.0)
    inputs = [t0_fp4]

    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[1024, 16],
            contiguity=[True, True],
            dtype=DataType.Float4_e2m1fn_x2,
            is_cpu=False,
        )
        T1 = fd.ops.cast(T0, DataType.Float)
        T2 = fd.ops.relu(T1)
        fd.add_output(T2)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    ref = fp4_to_fp32(unpack_fp4(t0_fp4.view(torch.uint8))).relu()
    nvfuser_direct_test.assertEqual(out[0], ref)

def test_broadcast_in_dim_no_redundant_set(nvfuser_direct_test):
    """
    Test that broadcast_in_dim doesn't introduce redundant Set operations
    when all input dimensions are in broadcast_dims (i.e., no actual broadcast).

    This verifies the fix for the issue where broadcast_in_dim would create
    a redundant float-to-float cast operation via Set when the input already
    had the correct shape.
    """

    def fusion_with_broadcast_in_dim(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # broadcast_in_dim with broadcast_dims=[0, 1] means no new dims are added
        t2 = fd.ops.broadcast_in_dim(t0, t1.shape(), [0, 1])
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)


    # we'll just ignore the scaling factor, since we only want to test basic fp4 support
    t0_fp4, _ = pytorch_nvfp4_quantize(t0, 1.0)
    inputs = [t0_fp4]

    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[1024, 16],
            contiguity=[True, True],
            dtype=DataType.Float4_e2m1fn_x2,
            is_cpu=False,
        )
        T1 = fd.ops.cast(T0, DataType.Float)
        T2 = fd.ops.relu(T1)
        fd.add_output(T2)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    ref = fp4_to_fp32(unpack_fp4(t0_fp4.view(torch.uint8))).relu()
    nvfuser_direct_test.assertEqual(out[0], ref)

def test_broadcast_in_dim_no_redundant_set(nvfuser_direct_test):
    """
    Test that broadcast_in_dim doesn't introduce redundant Set operations
    when all input dimensions are in broadcast_dims (i.e., no actual broadcast).

    This verifies the fix for the issue where broadcast_in_dim would create
    a redundant float-to-float cast operation via Set when the input already
    had the correct shape.
    """

    def fusion_with_broadcast_in_dim(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # broadcast_in_dim with broadcast_dims=[0, 1] means no new dims are added
        t2 = fd.ops.broadcast_in_dim(t0, t1.shape(), [0, 1])
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    def fusion_with_expand(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # Direct expand without broadcast_in_dim
        t2 = fd.ops.expand(t0, t1.shape())
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    with FusionDefinition() as fd_bid:
        fusion_with_broadcast_in_dim(fd_bid)

    with FusionDefinition() as fd_exp:
        fusion_with_expand(fd_exp)

    # Check that the broadcast_in_dim fusion doesn't have a redundant Set operation
    # by comparing the IR string representations - they should be identical since
    # broadcast is a no-op in this case
    assert str(fd_bid) == str(fd_exp)

def test_expanded_to_size_one(nvfuser_direct_test):
    """
    Test expanded to size one, which was failing a false positive assert
    """
    inputs = [
        torch.randint(0, 10, (1, 1), device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, None], dtype=DataType.Int
        )
        fd.add_output(tv0)

    with FusionDefinition() as fd:
        fusion_func(fd)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)


    def fusion_with_expand(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # Direct expand without broadcast_in_dim
        t2 = fd.ops.expand(t0, t1.shape())
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    with FusionDefinition() as fd_bid:
        fusion_with_broadcast_in_dim(fd_bid)

    with FusionDefinition() as fd_exp:
        fusion_with_expand(fd_exp)

    # Check that the broadcast_in_dim fusion doesn't have a redundant Set operation
    # by comparing the IR string representations - they should be identical since
    # broadcast is a no-op in this case
    assert str(fd_bid) == str(fd_exp)

def test_expanded_to_size_one(nvfuser_direct_test):
    """
    Test expanded to size one, which was failing a false positive assert
    """
    inputs = [
        torch.randint(0, 10, (1, 1), device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, None], dtype=DataType.Int
        )
        fd.add_output(tv0)

    with FusionDefinition() as fd:
        fusion_func(fd)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(out[0], inputs[0])

def test_issue4888(nvfuser_direct_test):
    # https://github.com/NVIDIA/Fuser/issues/4888
    def nvfuser_fusion_id2(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.Bool,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.define_tensor(


    """
    Test expanded to size one, which was failing a false positive assert
    """
    inputs = [
        torch.randint(0, 10, (1, 1), device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, None], dtype=DataType.Int
        )
        fd.add_output(tv0)

    with FusionDefinition() as fd:
        fusion_func(fd)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(out[0], inputs[0])

def test_issue4888(nvfuser_direct_test):
    # https://github.com/NVIDIA/Fuser/issues/4888
    def nvfuser_fusion_id2(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.Bool,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.Bool,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T3 = fd.define_tensor(
            shape=[1, 32, 4096, 4096],
            contiguity=[None, True, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[3, 2, 1, 0],
        )
        T4 = fd.ops.cast(T0, dtype=DataType.Float)
        T5 = fd.ops.bitwise_or(T1, T2)
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T1)
        T7 = fd.ops.cast(T6, dtype=DataType.Float)
        T8 = fd.ops.mul(T4, T7)
        T9 = fd.ops.cast(T8, dtype=DataType.BFloat16)
        T10 = fd.ops.set(T9)
        fd.add_output(T10, T0)
        T15 = fd.ops.broadcast_in_dim(T10, shape=[1, 4096, 4097], broadcast_dims=[1, 2])
        T21 = fd.ops.broadcast_in_dim(
            T15, shape=[1, 1, 4096, 4097], broadcast_dims=[0, 2, 3]
        )
        T27 = fd.ops.broadcast_in_dim(
            T21, shape=[1, 1, 4096, 4097], broadcast_dims=[0, 1, 2, 3]
        )
        T43 = fd.ops.slice(
            T27,
            start_indices=[0, 0, 0, 0],
            end_indices=[1, 1, 4096, 4096],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        T49 = fd.ops.broadcast_in_dim(
            T43, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T50 = fd.ops.cast(T49, dtype=DataType.Float)
        T51 = fd.ops.cast(T3, dtype=DataType.Float)
        S52 = fd.define_scalar(0.0883883, dtype=DataType.Double)
        T53 = fd.ops.mul(T51, S52)
        T54 = fd.ops.add(T53, T50)
        T55 = fd.ops.max(T54, dims=[3], keepdim=False, dtype=DataType.Null)
        T61 = fd.ops.broadcast_in_dim(
            T55, shape=[1, 32, 4096, 1], broadcast_dims=[0, 1, 2]
        )
        T67 = fd.ops.broadcast_in_dim(
            T61, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T68 = fd.ops.sub(T54, T67)
        T69 = fd.ops.exp(T68)
        T70 = fd.ops.sum(T69, dims=[3], keepdim=False, dtype=DataType.Null)
        T76 = fd.ops.broadcast_in_dim(
            T70, shape=[1, 32, 4096, 1], broadcast_dims=[0, 1, 2]
        )
        T82 = fd.ops.broadcast_in_dim(
            T76, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T83 = fd.ops.reciprocal(T82)
        T84 = fd.ops.mul(T69, T83)
        T85 = fd.ops.cast(T84, dtype=DataType.BFloat16)
        fd.add_output(T49)
        fd.add_output(T84)
        fd.add_output(T85)

    inputs = [
        torch.testing.make_tensor((4096, 4097), dtype=torch.bfloat16, device="cuda:0"),
        torch.testing.make_tensor((4096, 4097), dtype=torch.bool, device="cuda:0"),
        torch.testing.make_tensor((4096, 4097), dtype=torch.bool, device="cuda:0"),
        torch.testing.make_tensor(
            (1, 32, 4096, 4096), dtype=torch.bfloat16, device="cuda:0"
        ),
    ]
    nvfuser_direct_test.exec_nvfuser(nvfuser_fusion_id2, inputs)


        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        expected_output = torch.zeros_like(current_input)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(None, dtype=DataType.Int)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_func, [current_input, num_bits]
        )
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_output)

def test_all_dim_var_mean(nvfuser_direct_test):
    inputs = [torch.randn(2, 2, 2, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [0, 1, 2], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    list_of_test_cases = [0, 1]
    for correction in list_of_test_cases:
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [0, 1, 2], bool(correction))
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_validate_precomputed_values(nvfuser_direct_test):
    # This test is from legacy test_nan.py
    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,


fn

        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])

        t0h = fd.ops.cast(t0, DataType.Half)
        t1h = fd.ops.cast(t1, DataType.Half)
        t2 = fd.ops.add(t0h, t1h)
        t3 = fd.ops.relu(t2)
        t4 = fd.ops.cast(t3, DataType.Half)

        fd.add_output(t4)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.relu(inputs[0].to(torch.half) + inputs[1].to(torch.half))
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_cast_fp8(nvfuser_direct_test):
    def fn(in_type, out_type):
        inputs = [
            torch.randn([5, 5], device="cuda").to(in_type),
        ]

        def fusion_func(fd: FusionDefinition) -> None:
            T0 = fd.from_pytorch(inputs[0])
            T1 = fd.ops.cast(T0, dtype=torch_dtype_to_nvfuser_dtype(out_type))
            fd.add_output(T1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = inputs[0].to(out_type)
        if in_type == torch.float8_e8m0fnu or out_type == torch.float8_e8m0fnu:
            # Eager mode uses manual bit manipulation, and nvFuser uses
            # hardware instructions. Unfortunately, these implementations
            # do not match exactly. e8m0 can only represent 2^x, so we are
            # asserting that the x of the two results are off by at most 1.
            nvf_out_fp32 = nvf_out[0].to(torch.float32)
            eager_out_fp32 = eager_out.to(torch.float32)
            rel_err = eager_out_fp32.div(nvf_out_fp32).max().item()
            nvfuser_direct_test.assertTrue(rel_err <= 2 and rel_err >= 0.5)
        else:
            nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    for type0 in [torch.double, torch.float32, torch.float16, torch.bfloat16]:
        type1_list = [torch.float8_e4m3fn, torch.float8_e5m2]
        if not is_pre_blackwell():
            type1_list.append(torch.float8_e8m0fnu)
        for type1 in type1_list:
            fn(type0, type1)
            fn(type1, type0)

def test_promote_to_double(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda", dtype=torch.float16),
        torch.randn(2, 4, device="cuda", dtype=torch.float64),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])


        n_shape = fd.define_vector(2)

        t1 = fd.ops.reshape(t0, n_shape)
        t2 = fd.ops.sum(t1, dims=[0])

        fd.add_output(t2)

    eager_out = torch.sum(inputs_with_list[0].reshape(new_shape), dim=0)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs_with_list)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    inputs_with_tuple = [tensor, tuple(new_shape)]
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs_with_tuple, new_fusion_expected=False
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_dynamic_reshape(nvfuser_direct_test):
    def dynamic_reshape(fd: FusionDefinition) -> None:
        x = fd.define_tensor([-1, -1], [True, True])
        d0 = fd.ops.size(x, 0)
        d1 = fd.define_scalar(dtype=DataType.Int32)
        d2 = fd.define_scalar(dtype=DataType.Int32)
        y = fd.ops.reshape(x, [d0, d1, d2])
        fd.add_output(y)

    x = torch.rand(3, 4, device="cuda")
    ys, _ = nvfuser_direct_test.exec_nvfuser(dynamic_reshape, [x, 2, 2])
    nvfuser_direct_test.assertEqual(len(ys), 1)
    y = ys[0]

    nvfuser_direct_test.assertEqual(y.shape, torch.Size([3, 2, 2]))
    nvfuser_direct_test.assertEqual(x.flatten(), y.flatten())

def test_reshape_dynamic(nvfuser_direct_test):
    inputs = [
        32,
        torch.randn((192,), dtype=torch.float32, device="cuda:0").as_strided(
            (4, 8, 6), (48, 6, 1)
        ),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        S0 = fd.define_scalar(None, dtype=DataType.Int)


    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.cummin(t0, dim)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

        eager_out = torch.cummin(inputs[0], dim)
        nvfuser_direct_test.assertEqual(eager_out[0], nvf_out[0])

    test_fn(0)
    test_fn(1)

def test_cummax(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.cummax(t0, dim)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

        eager_out = torch.cummax(inputs[0], dim)
        nvfuser_direct_test.assertEqual(eager_out[0], nvf_out[0])

    test_fn(0)
    test_fn(1)

def test_where(nvfuser_direct_test):
    # nvfuser_where is a decorator function. It takes the input arguments
    # and creates a function that builds a FusionDefinition.
    def nvfuser_where(pred, a, b):
        def fusion_func(fd: FusionDefinition):
            nv_pred = fd.define_tensor(
                sizes=pred.shape, strides=pred.stride(), dtype=DataType.Bool
            )
            nv_a = fd.define_tensor(
                sizes=a.shape,
                strides=a.stride(),
                dtype=torch_dtype_to_nvfuser_dtype(a.dtype),
            )
            nv_b = fd.define_tensor(
                sizes=b.shape,


    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.cummax(t0, dim)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

        eager_out = torch.cummax(inputs[0], dim)
        nvfuser_direct_test.assertEqual(eager_out[0], nvf_out[0])

    test_fn(0)
    test_fn(1)

def test_where(nvfuser_direct_test):
    # nvfuser_where is a decorator function. It takes the input arguments
    # and creates a function that builds a FusionDefinition.
    def nvfuser_where(pred, a, b):
        def fusion_func(fd: FusionDefinition):
            nv_pred = fd.define_tensor(
                sizes=pred.shape, strides=pred.stride(), dtype=DataType.Bool
            )
            nv_a = fd.define_tensor(
                sizes=a.shape,
                strides=a.stride(),
                dtype=torch_dtype_to_nvfuser_dtype(a.dtype),
            )
            nv_b = fd.define_tensor(
                sizes=b.shape,
                strides=b.stride(),
                dtype=torch_dtype_to_nvfuser_dtype(b.dtype),
            )
            result = fd.ops.where(nv_pred, nv_a, nv_b)
            fd.add_output(result)

        return fusion_func

    # get list of dtypes to test with
    list_of_dtype = [torch.float16, torch.float32]
    if not is_pre_ampere():
        list_of_dtype.append(torch.bfloat16)

    pred = torch.testing.make_tensor((5,), device="cuda", dtype=torch.bool)
    for atype, btype in itertools.product(list_of_dtype, list_of_dtype):
        a = torch.randn((5,), device="cuda", dtype=atype)
        b = torch.randn((5,), device="cuda", dtype=btype)
        fusion_func = nvfuser_where(pred, a, b)
        nv_result, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [pred, a, b])
        torch_result = torch.where(pred, a, b)
        nvfuser_direct_test.assertEqual(nv_result[0], torch_result)

def test_where_dtypes(nvfuser_direct_test):
    inputs = [
        torch.arange(2, device="cuda").type(torch.bool),
    ]


reshape

        torch.testing.make_tensor(1024, device="cuda", dtype=torch.long),
        torch.testing.make_tensor(1024, device="cuda", dtype=torch.long),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        t3 = fd.ops.truediv(t0, t1)
        fd.add_output(t2)
        fd.add_output(t3)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(
        nvf_out[0], torch.div(inputs[0], inputs[1], rounding_mode="trunc")
    )
    nvfuser_direct_test.assertEqual(nvf_out[1], torch.true_divide(inputs[0], inputs[1]))

def test_mark_alias_pass(nvfuser_direct_test):
    def reshape(fd: FusionDefinition) -> None:
        x = fd.define_tensor(
            [2, 3, 4], contiguity=[True, True, True], dtype=DataType.Float
        )
        y = fd.ops.reshape(x, [2, 12])
        fd.add_output(y)

    x = torch.rand(2, 3, 4, device="cuda")
    ys, _ = nvfuser_direct_test.exec_nvfuser(reshape, [x])
    nvfuser_direct_test.assertEqual(len(ys), 1)
    y = ys[0]

    nvfuser_direct_test.assertEqual(y.data_ptr(), x.data_ptr())

def test_misaligned_add(nvfuser_direct_test):
    inputs = [
        torch.ones(2**20 + 1, device="cuda")[1:],  # cannot vectorize
        torch.ones(2**20, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)


def test_real_imag(nvfuser_direct_test):
    for dtype in [torch.complex128, torch.complex64]:
        inputs = [
            torch.randn(5, dtype=dtype, device="cuda"),
        ]

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            fd.add_output(fd.ops.real(t0))
            fd.add_output(fd.ops.imag(t0))

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

        nvfuser_direct_test.assertEqual(torch.real(inputs[0]), nvf_out[0])
        nvfuser_direct_test.assertEqual(torch.imag(inputs[0]), nvf_out[1])

def test_reduction_complex_number(nvfuser_direct_test):
    def test_dtype(torch_dtype):
        inputs = [torch.randn(2, 32, device="cuda", dtype=torch_dtype)]

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.sum(t0, [-1], False, torch_dtype_to_nvfuser_dtype(torch_dtype))
            fd.add_output(t1)

        nvf_out1, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = torch.sum(inputs[0], dim=-1)
        nvfuser_direct_test.assertEqual(eager_out, nvf_out1[0])

    list_of_dtype = [torch.complex64, torch.complex128]
    for torch_dtype in list_of_dtype:
        test_dtype(torch_dtype)

def test_right_shift_arithmetic(nvfuser_direct_test):
    inputs = [torch.tensor([-2147483648, 1073741824], dtype=torch.int32, device="cuda")]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        c0 = fd.define_scalar(3)
        t1 = fd.ops.bitwise_right_shift(t0, c0)
        fd.add_output(t1)

    nvf_out1, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out1 = torch.bitwise_right_shift(inputs[0], 3)
    nvfuser_direct_test.assertEqual(eager_out1, nvf_out1[0])


        # >>> jax.lax.shift_right_logical(input.to('cpu').numpy(), 64)
        # Array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
        expected_output = torch.zeros_like(current_input)

        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(current_input)
            c0 = fd.define_scalar(None, dtype=DataType.Int)
            t1 = fd.ops.logical_right_shift(t0, c0)
            fd.add_output(t1)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_func, [current_input, num_bits]
        )
        nvfuser_direct_test.assertEqual(nvf_out[0], expected_output)

def test_all_dim_var_mean(nvfuser_direct_test):
    inputs = [torch.randn(2, 2, 2, device="cuda")]

    # use decorator to create fusion_func
    def fusion_decorator(correction):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            t1, t2 = fd.ops.var_mean(t0, [0, 1, 2], correction)
            fd.add_output(t1)
            fd.add_output(t2)

        return fusion_func

    list_of_test_cases = [0, 1]
    for correction in list_of_test_cases:
        fuser_result, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_decorator(correction), inputs
        )
        torch_result = torch.var_mean(inputs[0], [0, 1, 2], bool(correction))
        nvfuser_direct_test.assertEqual(fuser_result, torch_result)

def test_validate_precomputed_values(nvfuser_direct_test):
    # This test is from legacy test_nan.py
    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )


        )
        T1 = fd.ops.cast(T0, DataType.Float)
        T2 = fd.ops.relu(T1)
        fd.add_output(T2)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    ref = fp4_to_fp32(unpack_fp4(t0_fp4.view(torch.uint8))).relu()
    nvfuser_direct_test.assertEqual(out[0], ref)

def test_broadcast_in_dim_no_redundant_set(nvfuser_direct_test):
    """
    Test that broadcast_in_dim doesn't introduce redundant Set operations
    when all input dimensions are in broadcast_dims (i.e., no actual broadcast).

    This verifies the fix for the issue where broadcast_in_dim would create
    a redundant float-to-float cast operation via Set when the input already
    had the correct shape.
    """

    def fusion_with_broadcast_in_dim(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # broadcast_in_dim with broadcast_dims=[0, 1] means no new dims are added
        t2 = fd.ops.broadcast_in_dim(t0, t1.shape(), [0, 1])
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    def fusion_with_expand(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # Direct expand without broadcast_in_dim
        t2 = fd.ops.expand(t0, t1.shape())
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    with FusionDefinition() as fd_bid:
        fusion_with_broadcast_in_dim(fd_bid)

    with FusionDefinition() as fd_exp:
        fusion_with_expand(fd_exp)

    # Check that the broadcast_in_dim fusion doesn't have a redundant Set operation
    # by comparing the IR string representations - they should be identical since
    # broadcast is a no-op in this case
    assert str(fd_bid) == str(fd_exp)


def test_broadcast_in_dim_no_redundant_set(nvfuser_direct_test):
    """
    Test that broadcast_in_dim doesn't introduce redundant Set operations
    when all input dimensions are in broadcast_dims (i.e., no actual broadcast).

    This verifies the fix for the issue where broadcast_in_dim would create
    a redundant float-to-float cast operation via Set when the input already
    had the correct shape.
    """

    def fusion_with_broadcast_in_dim(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # broadcast_in_dim with broadcast_dims=[0, 1] means no new dims are added
        t2 = fd.ops.broadcast_in_dim(t0, t1.shape(), [0, 1])
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    def fusion_with_expand(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[1, -1], contiguity=[None, True])
        t1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True])
        # Direct expand without broadcast_in_dim
        t2 = fd.ops.expand(t0, t1.shape())
        t3 = fd.ops.add(t2, t1)
        fd.add_output(t3)

    with FusionDefinition() as fd_bid:
        fusion_with_broadcast_in_dim(fd_bid)

    with FusionDefinition() as fd_exp:
        fusion_with_expand(fd_exp)

    # Check that the broadcast_in_dim fusion doesn't have a redundant Set operation
    # by comparing the IR string representations - they should be identical since
    # broadcast is a no-op in this case
    assert str(fd_bid) == str(fd_exp)

def test_expanded_to_size_one(nvfuser_direct_test):
    """
    Test expanded to size one, which was failing a false positive assert
    """
    inputs = [
        torch.randint(0, 10, (1, 1), device="cuda"),
    ]


    """
    inputs = [
        torch.randint(0, 10, (1, 1), device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, None], dtype=DataType.Int
        )
        fd.add_output(tv0)

    with FusionDefinition() as fd:
        fusion_func(fd)

    out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    nvfuser_direct_test.assertEqual(out[0], inputs[0])

def test_issue4888(nvfuser_direct_test):
    # https://github.com/NVIDIA/Fuser/issues/4888
    def nvfuser_fusion_id2(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.Bool,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.define_tensor(
            shape=[4096, 4097],
            contiguity=[True, True],
            dtype=DataType.Bool,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T3 = fd.define_tensor(
            shape=[1, 32, 4096, 4096],
            contiguity=[None, True, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[3, 2, 1, 0],
        )
        T4 = fd.ops.cast(T0, dtype=DataType.Float)
        T5 = fd.ops.bitwise_or(T1, T2)
        T6 = fd.ops.set(T5)
        fd.add_output(T6, T1)
        T7 = fd.ops.cast(T6, dtype=DataType.Float)
        T8 = fd.ops.mul(T4, T7)
        T9 = fd.ops.cast(T8, dtype=DataType.BFloat16)
        T10 = fd.ops.set(T9)
        fd.add_output(T10, T0)
        T15 = fd.ops.broadcast_in_dim(T10, shape=[1, 4096, 4097], broadcast_dims=[1, 2])
        T21 = fd.ops.broadcast_in_dim(
            T15, shape=[1, 1, 4096, 4097], broadcast_dims=[0, 2, 3]
        )
        T27 = fd.ops.broadcast_in_dim(
            T21, shape=[1, 1, 4096, 4097], broadcast_dims=[0, 1, 2, 3]
        )
        T43 = fd.ops.slice(
            T27,
            start_indices=[0, 0, 0, 0],
            end_indices=[1, 1, 4096, 4096],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        T49 = fd.ops.broadcast_in_dim(
            T43, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T50 = fd.ops.cast(T49, dtype=DataType.Float)
        T51 = fd.ops.cast(T3, dtype=DataType.Float)
        S52 = fd.define_scalar(0.0883883, dtype=DataType.Double)
        T53 = fd.ops.mul(T51, S52)
        T54 = fd.ops.add(T53, T50)
        T55 = fd.ops.max(T54, dims=[3], keepdim=False, dtype=DataType.Null)
        T61 = fd.ops.broadcast_in_dim(
            T55, shape=[1, 32, 4096, 1], broadcast_dims=[0, 1, 2]
        )
        T67 = fd.ops.broadcast_in_dim(
            T61, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T68 = fd.ops.sub(T54, T67)
        T69 = fd.ops.exp(T68)
        T70 = fd.ops.sum(T69, dims=[3], keepdim=False, dtype=DataType.Null)
        T76 = fd.ops.broadcast_in_dim(
            T70, shape=[1, 32, 4096, 1], broadcast_dims=[0, 1, 2]
        )
        T82 = fd.ops.broadcast_in_dim(
            T76, shape=[1, 32, 4096, 4096], broadcast_dims=[0, 1, 2, 3]
        )
        T83 = fd.ops.reciprocal(T82)
        T84 = fd.ops.mul(T69, T83)
        T85 = fd.ops.cast(T84, dtype=DataType.BFloat16)
        fd.add_output(T49)
        fd.add_output(T84)
        fd.add_output(T85)

    inputs = [
        torch.testing.make_tensor((4096, 4097), dtype=torch.bfloat16, device="cuda:0"),
        torch.testing.make_tensor((4096, 4097), dtype=torch.bool, device="cuda:0"),
        torch.testing.make_tensor((4096, 4097), dtype=torch.bool, device="cuda:0"),
        torch.testing.make_tensor(
            (1, 32, 4096, 4096), dtype=torch.bfloat16, device="cuda:0"
        ),
    ]
    nvfuser_direct_test.exec_nvfuser(nvfuser_fusion_id2, inputs)


        .min()
        .cpu()
        .float()
        .isclose(torch.tensor(lo), rtol=1e-2, atol=1e-2)
        .item()
    )
    nvfuser_direct_test.assertTrue(
        nvf_out[0]
        .max()
        .cpu()
        .float()
        .isclose(torch.tensor(hi), rtol=1e-2, atol=1e-2)
        .item()
    )

def test_random_distinct_values(nvfuser_direct_test):
    dtypes = [DataType.Double, DataType.Float, DataType.Half, DataType.BFloat16]
    for dtype, rand_op_name in itertools.product(dtypes, ["uniform", "normal"]):

        def fusion_fn(fd: FusionDefinition):
            # generate 4 values and check that they are all distinct
            rand_op = getattr(fd.ops, rand_op_name)
            S0 = fd.define_scalar(0.00000, dtype=DataType.Double)
            S1 = fd.define_scalar(1.00000, dtype=DataType.Double)
            output = rand_op(S0, S1, shape=[2, 2], dtype=dtype)
            fd.add_output(output)

        with FusionDefinition() as fd:
            fusion_fn(fd)

        for i in range(100):
            output = fd.execute([])[0]

            # Rarely we might have a pair of matching lower precision
            # samples. However, it is extremely rare that we would have a
            # set of three matching elements in only 100 repeats unless we
            # have a bug.

            match = output.flatten().unsqueeze(0) == output.flatten().unsqueeze(1)
            match_pairs = (
                match ^ torch.eye(4, dtype=torch.bool, device="cuda")
            ).sum() // 2

            assert match_pairs.item() < 3, f"At least three entries match in {output}"


        fd.add_output(t4)
        fd.add_output(t5)
        fd.add_output(t6)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = inputs[0] + inputs[1]
    nvfuser_direct_test.assertEqual(nvf_out[0], inputs[0] + inputs[1])
    nvfuser_direct_test.assertEqual(nvf_out[1], inputs[2] + 3.0)
    nvfuser_direct_test.assertEqual(nvf_out[2], inputs[3] * 3.0)

def test_output_stride_order(nvfuser_direct_test):
    inputs = [
        torch.arange(0, 24).reshape(2, 3, 4).cuda().float(),
    ]
    eager_out = inputs[0] + 3.0

    for perm in itertools.permutations(range(3), 3):
        # testing stride_order in set
        def fusion_set_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            c0 = fd.define_scalar(3.0)
            t1 = fd.ops.add(t0, c0)
            t2 = fd.ops.stride_order(t1, perm)
            fd.add_output(t2)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_set_func, inputs)
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

        nvf_stride = nvf_out[0].stride()
        verify_stride_order(nvf_stride, perm)

def test_output_stride_order_with_reduction(nvfuser_direct_test):
    inputs = [torch.randn(2, 3, 4, 5, device="cuda", dtype=torch.float)]

    for stride_order in itertools.permutations(range(3), 3):

        def fusion_stride_order_op(fd: FusionDefinition) -> None:
            T0 = fd.from_pytorch(inputs[0])
            T1 = fd.ops.sum(T0, dims=[2])
            T2 = fd.ops.stride_order(T1, stride_order)
            fd.add_output(T2)

        with FusionDefinition() as fd:


        # testing stride_order in set
        def fusion_set_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            c0 = fd.define_scalar(3.0)
            t1 = fd.ops.add(t0, c0)
            t2 = fd.ops.stride_order(t1, perm)
            fd.add_output(t2)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_set_func, inputs)
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

        nvf_stride = nvf_out[0].stride()
        verify_stride_order(nvf_stride, perm)

def test_output_stride_order_with_reduction(nvfuser_direct_test):
    inputs = [torch.randn(2, 3, 4, 5, device="cuda", dtype=torch.float)]

    for stride_order in itertools.permutations(range(3), 3):

        def fusion_stride_order_op(fd: FusionDefinition) -> None:
            T0 = fd.from_pytorch(inputs[0])
            T1 = fd.ops.sum(T0, dims=[2])
            T2 = fd.ops.stride_order(T1, stride_order)
            fd.add_output(T2)

        with FusionDefinition() as fd:
            fusion_stride_order_op(fd)

        out = fd.execute(inputs)[0]
        verify_stride_order(out.stride(), stride_order)

def test_triu(nvfuser_direct_test):
    inputs = [
        torch.randn(4, 16, device="cuda", dtype=torch.float16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.ops.triu(t0, -1)
        fd.add_output(t1)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out0 = torch.triu(inputs[0], -1)