    from torch.nn.attention import SDPBackend, sdpa_kernel

    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            fusion_func,
            inputs,
        )
    # Ignoring size-1 dimensions, `q @ k^T / sqrt(e)` generates a `l`x`s`
    # matrix full of `sqrt(e)`s.  Therefore, the logsumexp of each row is
    # expected to be log(exp(sqrt(e)) * s) = log(s) + sqrt(e).
    torch.testing.assert_close(
        nvf_out[0].cpu(), torch.full((n, h, l), math.log(s) + e**0.5)
    )

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
def test_sdpa_fwd(nvfuser_direct_test):
    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        v = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        dropout_p, is_causal, scale = None, None, None
        if has_dropout:
            dropout_p = fd.define_scalar(value=None, dtype=DataType.Double)
        if has_causal:
            is_causal = fd.define_scalar(value=None, dtype=DataType.Bool)
        if has_scale:
            scale = fd.define_scalar(value=None, dtype=DataType.Double)
        attn, *_ = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal, scale=scale
        )
        fd.add_output(attn)

    N, H, L, S, E = 4, 8, 16, 16, 8
    qkv = [
        torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda"),
        torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda"),
        torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda"),
    ]

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]
    # TODO: Try to move this to pytest_ops.py. Currently, it does not work since the API between nvFuser and torch differs.
    for dropout_p, is_causal, scale in itertools.product(
        dropout_vals, is_causal_vals, scale_vals
    ):
        with nvfuser_direct_test.subTest(
            dropout_p=dropout_p, is_causal=is_causal, scale=scale
        ):
            from torch.nn.attention import SDPBackend, sdpa_kernel

            has_dropout = True if dropout_p is not None else False
            has_causal = True if is_causal is not None else False
            has_scale = True if scale is not None else False
            inputs = [*qkv]
            for param in [dropout_p, is_causal, scale]:
                if param is not None:
                    inputs.append(param)
            nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
                partial(
                    fusion_func,
                    has_dropout=has_dropout,
                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )

            # Torch does not accept NoneType dropout_p, is_causal.
            dropout_p = 0.0 if dropout_p is None else dropout_p
            is_causal = False if is_causal is None else is_causal

            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                torch.manual_seed(0)
                ref_out = F.scaled_dot_product_attention(
                    *qkv, dropout_p=dropout_p, is_causal=is_causal, scale=scale
                )
            torch.testing.assert_close(nvf_out[0], ref_out)

# Memory layout of query, key, value and output tensors.
class Layout(Enum):
    NHSE = auto()
    NSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("layout", [Layout.NHSE, Layout.NSHE])
def test_sdpa_fwd_bias_mask(nvfuser_direct_test, layout: Layout):
    match layout:
        case Layout.NHSE:
            stride_order = [3, 2, 1, 0]
        case Layout.NSHE:
            stride_order = [3, 1, 2, 0]


Layout

                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )

            # Torch does not accept NoneType dropout_p, is_causal.
            dropout_p = 0.0 if dropout_p is None else dropout_p
            is_causal = False if is_causal is None else is_causal

            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                torch.manual_seed(0)
                ref_out = F.scaled_dot_product_attention(
                    *qkv, dropout_p=dropout_p, is_causal=is_causal, scale=scale
                )
            torch.testing.assert_close(nvf_out[0], ref_out)

# Memory layout of query, key, value and output tensors.
class Layout(Enum):
    NHSE = auto()
    NSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("layout", [Layout.NHSE, Layout.NSHE])
def test_sdpa_fwd_bias_mask(nvfuser_direct_test, layout: Layout):
    match layout:
        case Layout.NHSE:
            stride_order = [3, 2, 1, 0]
        case Layout.NSHE:
            stride_order = [3, 1, 2, 0]

    with FusionDefinition() as fd:
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            stride_order=stride_order,


            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                torch.manual_seed(0)
                ref_out = F.scaled_dot_product_attention(
                    *qkv, dropout_p=dropout_p, is_causal=is_causal, scale=scale
                )
            torch.testing.assert_close(nvf_out[0], ref_out)

# Memory layout of query, key, value and output tensors.
class Layout(Enum):
    NHSE = auto()
    NSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("layout", [Layout.NHSE, Layout.NSHE])
def test_sdpa_fwd_bias_mask(nvfuser_direct_test, layout: Layout):
    match layout:
        case Layout.NHSE:
            stride_order = [3, 2, 1, 0]
        case Layout.NSHE:
            stride_order = [3, 1, 2, 0]

    with FusionDefinition() as fd:
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            stride_order=stride_order,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            stride_order=stride_order,
        )
        v = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            stride_order=stride_order,
        )
        bias = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
        )
        mask = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.Bool,
        )
        attn, *_ = fd.ops.sdpfa_fwd(q, k, v, bias=bias, mask=mask)
        fd.add_output(attn)

    N, H, L, S, E = 11, 7, 5, 3, 2
    match layout:
        case Layout.NHSE:
            q = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")
            k = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")
            v = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")
        case Layout.NSHE:
            q = torch.randn(
                (N, L, H, E), dtype=torch.bfloat16, device="cuda"
            ).transpose(1, 2)
            k = torch.randn(
                (N, S, H, E), dtype=torch.bfloat16, device="cuda"
            ).transpose(1, 2)
            v = torch.randn(
                (N, S, H, E), dtype=torch.bfloat16, device="cuda"
            ).transpose(1, 2)
    bias = torch.randn((N, H, L, S), dtype=torch.bfloat16, device="cuda")
    mask = torch.rand((N, H, L, S), device="cuda") > 0.3

    (out,) = fd.execute([q, k, v, bias, mask])

    attn_mask = (bias + torch.where(mask, 0.0, float("-inf"))).to(dtype=bias.dtype)
    ref_out = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask
    )
    torch.testing.assert_close(out, ref_out)
    verify_stride_order(out.stride(), stride_order)

def test_sdpa_bwd(nvfuser_direct_test):
    N, H, L, S, E = 4, 8, 16, 16, 8

    grad_output = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")
    q = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")
    k = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")
    v = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]

    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        grad_output = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,


            ).transpose(1, 2)
            k = torch.randn(
                (N, S, H, E), dtype=torch.bfloat16, device="cuda"
            ).transpose(1, 2)
            v = torch.randn(
                (N, S, H, E), dtype=torch.bfloat16, device="cuda"
            ).transpose(1, 2)
    bias = torch.randn((N, H, L, S), dtype=torch.bfloat16, device="cuda")
    mask = torch.rand((N, H, L, S), device="cuda") > 0.3

    (out,) = fd.execute([q, k, v, bias, mask])

    attn_mask = (bias + torch.where(mask, 0.0, float("-inf"))).to(dtype=bias.dtype)
    ref_out = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask
    )
    torch.testing.assert_close(out, ref_out)
    verify_stride_order(out.stride(), stride_order)

def test_sdpa_bwd(nvfuser_direct_test):
    N, H, L, S, E = 4, 8, 16, 16, 8

    grad_output = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")
    q = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")
    k = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")
    v = torch.randn((N, H, S, E), dtype=torch.bfloat16, device="cuda")

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]

    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        grad_output = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        v = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        output = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        logsumexp = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=True,
            dtype=DataType.Float,
            is_cpu=False,
        )
        philox_seed, philox_offset = define_sdpa_rng_state(fd)

        dropout_p, is_causal, scale = None, None, None
        if has_dropout:
            dropout_p = fd.define_scalar(value=None, dtype=DataType.Double)
        if has_causal:
            is_causal = fd.define_scalar(value=None, dtype=DataType.Bool)
        if has_scale:
            scale = fd.define_scalar(value=None, dtype=DataType.Double)

        grad_query, grad_key, grad_value = fd.ops.sdpfa_bwd(
            grad_output,
            q,
            k,
            v,
            output,
            logsumexp,
            dropout_p,
            is_causal,
            philox_seed,
            philox_offset,
            scale,
        )
        fd.add_output(grad_query)
        fd.add_output(grad_key)
        fd.add_output(grad_value)

    for dropout_p, is_causal, scale in itertools.product(
        dropout_vals, is_causal_vals, scale_vals
    ):
        with nvfuser_direct_test.subTest(
            dropout_p=dropout_p, is_causal=is_causal, scale=scale
        ):
            # Torch does not accept NoneType dropout_p, is_causal.
            at_dropout_p = 0.0 if dropout_p is None else dropout_p
            at_is_causal = False if is_causal is None else is_causal

            (
                output,
                logsumexp,
                cum_seq_q,
                cum_seq_k,
                query_seq_len,
                key_seq_len,
                philox_seed,
                philox_offset,
                _,
            ) = torch.ops.aten._scaled_dot_product_flash_attention(
                q,
                k,
                v,
                at_dropout_p,
                at_is_causal,
                return_debug_mask=False,
                scale=scale,
            )
            ref_grad = torch.ops.aten._scaled_dot_product_flash_attention_backward(
                grad_output,
                q,
                k,
                v,
                output,
                logsumexp,
                cum_seq_q,
                cum_seq_k,
                query_seq_len,
                key_seq_len,
                at_dropout_p,
                at_is_causal,
                philox_seed,
                philox_offset,
                scale=scale,
            )

            has_dropout = True if dropout_p is not None else False
            has_causal = True if is_causal is not None else False
            has_scale = True if scale is not None else False

            inputs = [
                grad_output,
                q,
                k,
                v,
                output,
                logsumexp,
                philox_seed,
                philox_offset,
            ]
            for param in [dropout_p, is_causal, scale]:
                if param is not None:
                    inputs.append(param)

            nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
                partial(
                    fusion_func,
                    has_dropout=has_dropout,
                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )
            torch.testing.assert_close(nvf_out[0], ref_grad[0])
            torch.testing.assert_close(nvf_out[1], ref_grad[1])
            torch.testing.assert_close(nvf_out[2], ref_grad[2])

def test_sdpa_fwd_bwd(nvfuser_direct_test):
    N, H, L, S, E = 4, 8, 16, 16, 8

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]

    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],


            ]
            for param in [dropout_p, is_causal, scale]:
                if param is not None:
                    inputs.append(param)

            nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
                partial(
                    fusion_func,
                    has_dropout=has_dropout,
                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )
            torch.testing.assert_close(nvf_out[0], ref_grad[0])
            torch.testing.assert_close(nvf_out[1], ref_grad[1])
            torch.testing.assert_close(nvf_out[2], ref_grad[2])

def test_sdpa_fwd_bwd(nvfuser_direct_test):
    N, H, L, S, E = 4, 8, 16, 16, 8

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]

    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        v = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        grad_out = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )

        dropout_p, is_causal, scale = None, None, None
        if has_dropout:
            dropout_p = fd.define_scalar(value=None, dtype=DataType.Double)
        if has_causal:
            is_causal = fd.define_scalar(value=None, dtype=DataType.Bool)
        if has_scale:
            scale = fd.define_scalar(value=None, dtype=DataType.Double)

        output, logsumexp, philox_seed, philox_offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal, scale=scale
        )
        grad_query, grad_key, grad_value = fd.ops.sdpfa_bwd(
            grad_out,
            q,
            k,
            v,
            output,
            logsumexp,
            dropout_p,
            is_causal,
            philox_seed,
            philox_offset,
            scale,
        )

        fd.add_output(output)
        fd.add_output(grad_query)
        fd.add_output(grad_key)
        fd.add_output(grad_value)

    for dropout_p, is_causal, scale in itertools.product(
        dropout_vals, is_causal_vals, scale_vals
    ):
        with nvfuser_direct_test.subTest(
            dropout_p=dropout_p, is_causal=is_causal, scale=scale
        ):
            from torch.nn.attention import SDPBackend, sdpa_kernel

            q = torch.randn(
                (N, H, L, E),
                dtype=torch.bfloat16,
                device="cuda",
                requires_grad=True,
            )
            k = torch.randn(
                (N, H, S, E),
                dtype=torch.bfloat16,
                device="cuda",
                requires_grad=True,
            )
            v = torch.randn(
                (N, H, S, E),
                dtype=torch.bfloat16,
                device="cuda",
                requires_grad=True,
            )
            grad_output = torch.randn((N, H, L, E), dtype=torch.bfloat16, device="cuda")

            has_dropout = True if dropout_p is not None else False
            has_causal = True if is_causal is not None else False
            has_scale = True if scale is not None else False

            inputs = [q, k, v, grad_output]
            for param in [dropout_p, is_causal, scale]:
                if param is not None:
                    inputs.append(param)

            # Torch does not accept NoneType dropout_p, is_causal.
            dropout_p = 0.0 if dropout_p is None else dropout_p
            is_causal = False if is_causal is None else is_causal

            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                torch.manual_seed(0)
                ref_out = F.scaled_dot_product_attention(
                    q, k, v, dropout_p=dropout_p, is_causal=is_causal, scale=scale
                )
                ref_out.backward(grad_output)

            nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
                partial(
                    fusion_func,
                    has_dropout=has_dropout,
                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )
            torch.testing.assert_close(nvf_out[0], ref_out)
            torch.testing.assert_close(nvf_out[1], q.grad)
            torch.testing.assert_close(nvf_out[2], k.grad)
            torch.testing.assert_close(nvf_out[3], v.grad)


                    fusion_func,
                    has_dropout=has_dropout,
                    has_causal=has_causal,
                    has_scale=has_scale,
                ),
                inputs,
                new_fusion_expected=None,
            )
            torch.testing.assert_close(nvf_out[0], ref_grad[0])
            torch.testing.assert_close(nvf_out[1], ref_grad[1])
            torch.testing.assert_close(nvf_out[2], ref_grad[2])

def test_sdpa_fwd_bwd(nvfuser_direct_test):
    N, H, L, S, E = 4, 8, 16, 16, 8

    dropout_vals = [None, 0.0, 0.2]
    is_causal_vals = [None, True, False]
    scale_vals = [None, 1 / E**0.5, 1e-3]

    def fusion_func(
        fd: FusionDefinition, has_dropout: bool, has_causal: bool, has_scale: bool
    ):
        q = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        k = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        v = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )
        grad_out = fd.define_tensor(
            shape=[-1, -1, -1, -1],
            contiguity=True,
            dtype=DataType.BFloat16,
            is_cpu=False,
        )

        dropout_p, is_causal, scale = None, None, None
        if has_dropout:
            dropout_p = fd.define_scalar(value=None, dtype=DataType.Double)
        if has_causal:
            is_causal = fd.define_scalar(value=None, dtype=DataType.Bool)
        if has_scale:
            scale = fd.define_scalar(value=None, dtype=DataType.Double)

        output, logsumexp, philox_seed, philox_offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal, scale=scale
        )
        grad_query, grad_key, grad_value = fd.ops.sdpfa_bwd(
            grad_out,
            q,
            k,
            v,
            output,
            logsumexp,
            dropout_p,
            is_causal,
            philox_seed,
            philox_offset,
            scale,
        )

        fd.add_output(output)
        fd.add_output(grad_query)
        fd.add_output(grad_key)
        fd.add_output(grad_value)

    for dropout_p, is_causal, scale in itertools.product(
        dropout_vals, is_causal_vals, scale_vals
    ):
        with nvfuser_direct_test.subTest(
            dropout_p=dropout_p, is_causal=is_causal, scale=scale
        ):
            from torch.nn.attention import SDPBackend, sdpa_kernel

            q = torch.randn(
                (N, H, L, E),
                dtype=torch.bfloat16,
                device="cuda",
                requires_grad=True,
            )
            k = torch.randn(
                (N, H, S, E),
                dtype=torch.bfloat16,
                device="cuda",
                requires_grad=True,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import torch._refs as refs
import torch._prims as prims

from nvfuser_direct import (
    FusionDefinition,
    DataType,
    compute_tensor_descriptor,
)
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

import pytest
import itertools
from python.direct_utils import (
    is_pre_ampere,
    is_pre_hopper,
    is_pre_blackwell,
    verify_stride_order,
)

from python.direct_utils.narrow_precision import (
    pytorch_nvfp4_quantize,
    fp4_to_fp32,
    unpack_fp4,
)

def test_basic(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda"),
        torch.ones(2, 4, 8, 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)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        fd.add_output(t4)

    # Check that keepdim argument is not in fd_str
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.add(tv0, tv1)
    tv3 = fd.ops.mul(tv2, 3.00000)
    tv4 = fd.ops.sum(tv3, dims=[2], dtype=DataType.Float)
    fd.add_output(tv4)"""

    # t0 and t1 are ones(2, 4, 8) tensors.
    # t2 = t0 + t1 = twos(2, 4, 8)
    # t3 = t2 * 3.0 = sixes(2,4,8)
    # t4 = sum(t3, dim=-1) = forty-eights(2, 4)
    # The expected output is a tensor of 48's.
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3.0, dim=-1)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_basic_fp16(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
    ]

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

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        t5 = fd.ops.cast(t4, DataType.Half)
        fd.add_output(t5)

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

def test_define_contiguous_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([2, 3], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(2, 3, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_noncontiguous_tensor(nvfuser_direct_test):
    in_tensor = torch.randn(8, device="cuda").as_strided([2, 3], [4, 1])

    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_broadcast_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, 1, device="cuda").as_strided([1, 2, 1], [0, 1, 0])
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_tensor_contiguity_with_stride_order(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2], contiguity=True, stride_order=[0, 1])
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_cast_scalar(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
    ]

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

        t2 = fd.ops.add(t0, t1)
        c1 = fd.ops.cast(c0, DataType.Int32)
        t3 = fd.ops.mul(t2, c1)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Int32)

        fd.add_output(t4)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3, dim=-1, dtype=torch.int32)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_cast_double_to_half(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda", dtype=torch.float64),
        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])

        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])

        t2 = fd.ops.add(t0, t1)
        t5 = fd.ops.relu(t2)

        fd.add_output(t5)

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

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

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

        t0_b = fd.ops.broadcast(t0, [True, False, True])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_implicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(3, device="cuda"),
        torch.randn(2, 3, 4, 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, [2, 3, 4], [1])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_explicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(1, 1, 4, device="cuda"),
        torch.randn(2, 3, 4, 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, inputs[1].size(), [0, 1, 2])
        t2 = fd.ops.add(t0_b, t1)

        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, 1, 2]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_broadcast_mixing(nvfuser_direct_test):
    inputs = [
        torch.randn(3, 1, device="cuda"),
        torch.randn(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, [3, 3], [0])
        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], [3, 3], [0]))
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_dynamic_bcast_in_dim(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, True], dtype=DataType.Float, is_cpu=False
        )
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.broadcast_in_dim(tv0, (c4, c5), (0, 1))
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    t0 = torch.randn([1, 10], dtype=torch.float32, device="cuda:0")
    c4 = 5
    c5 = 10
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [t0, c4, c5, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_dynamic_full(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        c6 = fd.define_scalar(None, dtype=DataType.Float)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.full((c4, c5), c6, DataType.Float)
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    c4 = 5
    c5 = 10
    c6 = 2.5
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [c4, c5, c6, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_tensor_ndim(nvfuser_direct_test):
    shape = [2 for i in range(12)]
    new_shape = shape[:9]
    new_shape.append(8)

    inputs = [torch.randn(shape, device="cuda"), new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        n_shape = fd.define_vector(10)

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

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_execute_with_tuple_and_list(nvfuser_direct_test):
    shape = [2, 3, 4]
    new_shape = [6, 4]

    tensor = torch.randn(shape, device="cuda")
    inputs_with_list = [tensor, new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs_with_list[0])
        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)
        T1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        S2 = fd.define_scalar(1, dtype=DataType.Int)
        S3 = fd.ops.mul(S2, S0)
        S4 = fd.ops.signbit(S3)
        S5 = fd.define_scalar(False, dtype=DataType.Bool)
        S6 = fd.ops.ne(S4, S5)
        S7 = fd.define_scalar(192, dtype=DataType.Int)
        S8 = fd.ops.fmod(S7, S3)
        S9 = fd.ops.cast(S8, dtype=DataType.Int)
        S10 = fd.define_scalar(0, dtype=DataType.Int)
        S11 = fd.ops.ne(S9, S10)
        S12 = fd.ops.bitwise_and(S6, S11)
        S13 = fd.define_scalar(192, dtype=DataType.Int)
        S14 = fd.ops.reciprocal(S3)
        S15 = fd.ops.mul(S13, S14)
        S16 = fd.ops.cast(S12, dtype=DataType.Int)
        S17 = fd.ops.sub(S15, S16)
        T19 = fd.ops.reshape(T1, new_shape=[S0, S17])
        T20 = fd.ops.sum(T19, dims=[1], keepdim=False, dtype=DataType.Null)
        fd.add_output(T20)

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

# Test empty symbolic tensors can be reshaped
# See https://github.com/NVIDIA/Fuser/issues/2362
def test_empty_reshape(nvfuser_direct_test):
    inputs = [torch.randint(0, 10, (0, 1, 2, 3, 4), dtype=torch.int64, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, 1, -1, -1, -1],
            contiguity=[False, None, True, True, True],
            dtype=DataType.Int,
            is_cpu=False,
            stride_order=[4, 3, 2, 1, 0],
        )
        S2 = fd.define_scalar(5, dtype=DataType.Int)
        S3 = fd.define_scalar(0, dtype=DataType.Int)
        T5 = fd.ops.reshape(T0, new_shape=[S2, S3])
        fd.add_output(T5)

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

def test_squeeze(nvfuser_direct_test):
    t0_sizes = [4]
    t1_sizes = [1, 4, 1]
    t2_sizes = [2, 1, 4]
    inputs = [
        torch.randn(*t0_sizes, device="cuda"),
        torch.randn(*t1_sizes, device="cuda"),
        torch.randn(*t2_sizes, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[-1], contiguity=[True])
        t1 = fd.define_tensor(sizes=t1_sizes, strides=[4, 1, 1])
        t2 = fd.define_tensor(sizes=t2_sizes, strides=[4, 4, 1])
        t3 = fd.ops.squeeze(t1, [0, -1])
        t4 = fd.ops.squeeze(t2, [-2])
        t5 = fd.ops.sum(t4, [0])
        t6 = fd.ops.mul(t0, t3)
        t7 = fd.ops.mul(t6, t5)
        fd.add_output(t7)

    # Check that squeeze does not print default argument: squeeze_expanded
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[1, -1, 1], contiguity=[None, True, None], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.define_tensor(shape=[-1, 1, -1], contiguity=[True, None, True], dtype=DataType.Float, is_cpu=False)
    tv3 = fd.ops.squeeze(tv1, dims=[0, 2])
    tv6 = fd.ops.mul(tv0, tv3)
    tv4 = fd.ops.squeeze(tv2, dims=[1])
    tv5 = fd.ops.sum(tv4, dims=[0], dtype=DataType.Float)
    tv7 = fd.ops.mul(tv6, tv5)
    fd.add_output(tv7)"""

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

    v1 = torch.sum(inputs[1], [0, -1])
    v2 = torch.sum(inputs[2], [0, 1])
    eager_out = inputs[0] * v1 * v2
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

# Test that expanded dimensions can be reduced properly
# See https://github.com/NVIDIA/Fuser/issues/1678
def test_expanded_reduction(nvfuser_direct_test):
    inputs = [torch.tensor(1.0, device="cuda").as_strided((2, 3), (0, 0))]

    for keepdim in [False, True]:

        def fusion_func(fd: FusionDefinition) -> None:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[None, None],
                dtype=DataType.Float,
                is_cpu=False,
                stride_order=[1, 0],
            )
            T1 = fd.ops.sum(T0, dims=[0], keepdim=keepdim, dtype=DataType.Null)
            fd.add_output(T1)

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

        nvfuser_direct_test.assertEqual(
            nvf_out[0], inputs[0].sum(dim=0, keepdim=keepdim)
        )

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

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

        t0_b = fd.ops.expand(t0, inputs[1].size())
        t2 = fd.ops.mul(t0_b, t1)

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = inputs[0].expand(inputs[1].size()) * inputs[1]
    nvfuser_direct_test.assertEqual(eager
# ... truncated ...


    compute_tensor_descriptor,
)
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

import pytest
import itertools
from python.direct_utils import (
    is_pre_ampere,
    is_pre_hopper,
    is_pre_blackwell,
    verify_stride_order,
)

from python.direct_utils.narrow_precision import (
    pytorch_nvfp4_quantize,
    fp4_to_fp32,
    unpack_fp4,
)

def test_basic(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda"),
        torch.ones(2, 4, 8, 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)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        fd.add_output(t4)

    # Check that keepdim argument is not in fd_str
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.add(tv0, tv1)
    tv3 = fd.ops.mul(tv2, 3.00000)
    tv4 = fd.ops.sum(tv3, dims=[2], dtype=DataType.Float)
    fd.add_output(tv4)"""

    # t0 and t1 are ones(2, 4, 8) tensors.
    # t2 = t0 + t1 = twos(2, 4, 8)
    # t3 = t2 * 3.0 = sixes(2,4,8)
    # t4 = sum(t3, dim=-1) = forty-eights(2, 4)
    # The expected output is a tensor of 48's.
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3.0, dim=-1)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_basic_fp16(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
    ]

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

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        t5 = fd.ops.cast(t4, DataType.Half)
        fd.add_output(t5)


    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.add(tv0, tv1)
    tv3 = fd.ops.mul(tv2, 3.00000)
    tv4 = fd.ops.sum(tv3, dims=[2], dtype=DataType.Float)
    fd.add_output(tv4)"""

    # t0 and t1 are ones(2, 4, 8) tensors.
    # t2 = t0 + t1 = twos(2, 4, 8)
    # t3 = t2 * 3.0 = sixes(2,4,8)
    # t4 = sum(t3, dim=-1) = forty-eights(2, 4)
    # The expected output is a tensor of 48's.
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3.0, dim=-1)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_basic_fp16(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
    ]

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

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        t5 = fd.ops.cast(t4, DataType.Half)
        fd.add_output(t5)

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

def test_define_contiguous_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([2, 3], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(2, 3, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_noncontiguous_tensor(nvfuser_direct_test):
    in_tensor = torch.randn(8, device="cuda").as_strided([2, 3], [4, 1])

    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)


        torch.ones(2, 4, 8, device="cuda", dtype=torch.float16),
    ]

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

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Float)

        t5 = fd.ops.cast(t4, DataType.Half)
        fd.add_output(t5)

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

def test_define_contiguous_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([2, 3], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(2, 3, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_noncontiguous_tensor(nvfuser_direct_test):
    in_tensor = torch.randn(8, device="cuda").as_strided([2, 3], [4, 1])

    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_broadcast_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)


        t5 = fd.ops.cast(t4, DataType.Half)
        fd.add_output(t5)

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

def test_define_contiguous_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([2, 3], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(2, 3, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_noncontiguous_tensor(nvfuser_direct_test):
    in_tensor = torch.randn(8, device="cuda").as_strided([2, 3], [4, 1])

    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_broadcast_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, 1, device="cuda").as_strided([1, 2, 1], [0, 1, 0])
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_tensor_contiguity_with_stride_order(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2], contiguity=True, stride_order=[0, 1])
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, device="cuda")


        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(2, 3, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_noncontiguous_tensor(nvfuser_direct_test):
    in_tensor = torch.randn(8, device="cuda").as_strided([2, 3], [4, 1])

    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_broadcast_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, 1, device="cuda").as_strided([1, 2, 1], [0, 1, 0])
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_tensor_contiguity_with_stride_order(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2], contiguity=True, stride_order=[0, 1])
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_cast_scalar(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
    ]

    def fusion_func(fd: FusionDefinition):


    def fusion_func(fd: FusionDefinition):
        inp = fd.from_pytorch(in_tensor)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_broadcast_tensor(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, 1, device="cuda").as_strided([1, 2, 1], [0, 1, 0])
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_tensor_contiguity_with_stride_order(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2], contiguity=True, stride_order=[0, 1])
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_cast_scalar(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
    ]

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

        t2 = fd.ops.add(t0, t1)
        c1 = fd.ops.cast(c0, DataType.Int32)
        t3 = fd.ops.mul(t2, c1)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Int32)

        fd.add_output(t4)


        inp = fd.define_tensor([1, 2, 1], contiguity=True)
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, 1, device="cuda").as_strided([1, 2, 1], [0, 1, 0])
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_define_tensor_contiguity_with_stride_order(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition):
        inp = fd.define_tensor([1, 2], contiguity=True, stride_order=[0, 1])
        out = fd.ops.add(inp, inp)
        fd.add_output(out)

    in_tensor = torch.randn(1, 2, device="cuda")
    out_tensors, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, [in_tensor])
    nvfuser_direct_test.assertEqual(out_tensors[0], in_tensor * 2)

def test_cast_scalar(nvfuser_direct_test):
    inputs = [
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
    ]

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

        t2 = fd.ops.add(t0, t1)
        c1 = fd.ops.cast(c0, DataType.Int32)
        t3 = fd.ops.mul(t2, c1)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Int32)

        fd.add_output(t4)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3, dim=-1, dtype=torch.int32)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_cast_double_to_half(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda", dtype=torch.float64),
        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])

        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)


        torch.ones(2, 4, 8, device="cuda", dtype=torch.int32),
    ]

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

        t2 = fd.ops.add(t0, t1)
        c1 = fd.ops.cast(c0, DataType.Int32)
        t3 = fd.ops.mul(t2, c1)
        t4 = fd.ops.sum(t3, [-1], False, DataType.Int32)

        fd.add_output(t4)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum((inputs[0] + inputs[1]) * 3, dim=-1, dtype=torch.int32)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_cast_double_to_half(nvfuser_direct_test):
    inputs = [
        torch.randn(2, 4, device="cuda", dtype=torch.float64),
        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])

        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


    def fusion_func(fd: FusionDefinition):
        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])

        t2 = fd.ops.add(t0, t1)
        t5 = fd.ops.relu(t2)

        fd.add_output(t5)

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


            # 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])

        t2 = fd.ops.add(t0, t1)
        t5 = fd.ops.relu(t2)

        fd.add_output(t5)

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

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

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

        t0_b = fd.ops.broadcast(t0, [True, False, True])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]


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])

        t2 = fd.ops.add(t0, t1)
        t5 = fd.ops.relu(t2)

        fd.add_output(t5)

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

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

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

        t0_b = fd.ops.broadcast(t0, [True, False, True])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_implicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(3, device="cuda"),
        torch.randn(2, 3, 4, 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, [2, 3, 4], [1])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]


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

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

        t0_b = fd.ops.broadcast(t0, [True, False, True])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_implicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(3, device="cuda"),
        torch.randn(2, 3, 4, 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, [2, 3, 4], [1])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_explicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(1, 1, 4, device="cuda"),
        torch.randn(2, 3, 4, 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, inputs[1].size(), [0, 1, 2])
        t2 = fd.ops.add(t0_b, t1)

        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, 1, 2]), inputs[1]


        torch.randn(3, device="cuda"),
        torch.randn(2, 3, 4, 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, [2, 3, 4], [1])
        t2 = fd.ops.add(t0_b, t1)

        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(), [1]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_explicit_broadcast_input(nvfuser_direct_test):
    inputs = [
        torch.randn(1, 1, 4, device="cuda"),
        torch.randn(2, 3, 4, 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, inputs[1].size(), [0, 1, 2])
        t2 = fd.ops.add(t0_b, t1)

        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, 1, 2]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_broadcast_mixing(nvfuser_direct_test):
    inputs = [
        torch.randn(3, 1, device="cuda"),
        torch.randn(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, [3, 3], [0])
        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], [3, 3], [0]))
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])


        torch.randn(1, 1, 4, device="cuda"),
        torch.randn(2, 3, 4, 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, inputs[1].size(), [0, 1, 2])
        t2 = fd.ops.add(t0_b, t1)

        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, 1, 2]), inputs[1]
    )
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_broadcast_mixing(nvfuser_direct_test):
    inputs = [
        torch.randn(3, 1, device="cuda"),
        torch.randn(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, [3, 3], [0])
        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], [3, 3], [0]))
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_dynamic_bcast_in_dim(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, True], dtype=DataType.Float, is_cpu=False
        )
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.broadcast_in_dim(tv0, (c4, c5), (0, 1))
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    t0 = torch.randn([1, 10], dtype=torch.float32, device="cuda:0")
    c4 = 5
    c5 = 10
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")


def test_broadcast_mixing(nvfuser_direct_test):
    inputs = [
        torch.randn(3, 1, device="cuda"),
        torch.randn(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, [3, 3], [0])
        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], [3, 3], [0]))
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_dynamic_bcast_in_dim(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        tv0 = fd.define_tensor(
            shape=[1, -1], contiguity=[None, True], dtype=DataType.Float, is_cpu=False
        )
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.broadcast_in_dim(tv0, (c4, c5), (0, 1))
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    t0 = torch.randn([1, 10], dtype=torch.float32, device="cuda:0")
    c4 = 5
    c5 = 10
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [t0, c4, c5, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_dynamic_full(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        c6 = fd.define_scalar(None, dtype=DataType.Float)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.full((c4, c5), c6, DataType.Float)
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    c4 = 5
    c5 = 10
    c6 = 2.5
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [c4, c5, c6, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(


        )
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.broadcast_in_dim(tv0, (c4, c5), (0, 1))
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    t0 = torch.randn([1, 10], dtype=torch.float32, device="cuda:0")
    c4 = 5
    c5 = 10
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [t0, c4, c5, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_dynamic_full(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        c6 = fd.define_scalar(None, dtype=DataType.Float)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.full((c4, c5), c6, DataType.Float)
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    c4 = 5
    c5 = 10
    c6 = 2.5
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [c4, c5, c6, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_tensor_ndim(nvfuser_direct_test):
    shape = [2 for i in range(12)]
    new_shape = shape[:9]
    new_shape.append(8)

    inputs = [torch.randn(shape, device="cuda"), new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        n_shape = fd.define_vector(10)

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

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3)


        c4 = fd.define_scalar(None, dtype=DataType.Int)
        c5 = fd.define_scalar(None, dtype=DataType.Int)
        c6 = fd.define_scalar(None, dtype=DataType.Float)
        tv1 = fd.define_tensor(
            shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
        )
        tv3 = fd.ops.full((c4, c5), c6, DataType.Float)
        tv4 = fd.ops.add(tv3, tv1)
        fd.add_output(tv4)

    c4 = 5
    c5 = 10
    c6 = 2.5
    t1 = torch.randn([5, 10], dtype=torch.float32, device="cuda:0")
    inputs = [c4, c5, c6, t1]
    out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, validate_results=True
    )

def test_tensor_ndim(nvfuser_direct_test):
    shape = [2 for i in range(12)]
    new_shape = shape[:9]
    new_shape.append(8)

    inputs = [torch.randn(shape, device="cuda"), new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        n_shape = fd.define_vector(10)

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

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_execute_with_tuple_and_list(nvfuser_direct_test):
    shape = [2, 3, 4]
    new_shape = [6, 4]

    tensor = torch.randn(shape, device="cuda")
    inputs_with_list = [tensor, new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs_with_list[0])
        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)


    shape = [2 for i in range(12)]
    new_shape = shape[:9]
    new_shape.append(8)

    inputs = [torch.randn(shape, device="cuda"), new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        n_shape = fd.define_vector(10)

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

        fd.add_output(t2)

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3)
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_execute_with_tuple_and_list(nvfuser_direct_test):
    shape = [2, 3, 4]
    new_shape = [6, 4]

    tensor = torch.randn(shape, device="cuda")
    inputs_with_list = [tensor, new_shape]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs_with_list[0])
        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())


        t0 = fd.from_pytorch(inputs_with_list[0])
        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)
        T1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        S2 = fd.define_scalar(1, dtype=DataType.Int)


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)
        T1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        S2 = fd.define_scalar(1, dtype=DataType.Int)
        S3 = fd.ops.mul(S2, S0)
        S4 = fd.ops.signbit(S3)
        S5 = fd.define_scalar(False, dtype=DataType.Bool)
        S6 = fd.ops.ne(S4, S5)
        S7 = fd.define_scalar(192, dtype=DataType.Int)
        S8 = fd.ops.fmod(S7, S3)
        S9 = fd.ops.cast(S8, dtype=DataType.Int)
        S10 = fd.define_scalar(0, dtype=DataType.Int)
        S11 = fd.ops.ne(S9, S10)
        S12 = fd.ops.bitwise_and(S6, S11)
        S13 = fd.define_scalar(192, dtype=DataType.Int)
        S14 = fd.ops.reciprocal(S3)
        S15 = fd.ops.mul(S13, S14)
        S16 = fd.ops.cast(S12, dtype=DataType.Int)
        S17 = fd.ops.sub(S15, S16)
        T19 = fd.ops.reshape(T1, new_shape=[S0, S17])
        T20 = fd.ops.sum(T19, dims=[1], keepdim=False, dtype=DataType.Null)
        fd.add_output(T20)

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

# Test empty symbolic tensors can be reshaped
# See https://github.com/NVIDIA/Fuser/issues/2362
def test_empty_reshape(nvfuser_direct_test):
    inputs = [torch.randint(0, 10, (0, 1, 2, 3, 4), dtype=torch.int64, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, 1, -1, -1, -1],
            contiguity=[False, None, True, True, True],
            dtype=DataType.Int,
            is_cpu=False,
            stride_order=[4, 3, 2, 1, 0],
        )
        S2 = fd.define_scalar(5, dtype=DataType.Int)
        S3 = fd.define_scalar(0, dtype=DataType.Int)
        T5 = fd.ops.reshape(T0, new_shape=[S2, S3])
        fd.add_output(T5)


        S7 = fd.define_scalar(192, dtype=DataType.Int)
        S8 = fd.ops.fmod(S7, S3)
        S9 = fd.ops.cast(S8, dtype=DataType.Int)
        S10 = fd.define_scalar(0, dtype=DataType.Int)
        S11 = fd.ops.ne(S9, S10)
        S12 = fd.ops.bitwise_and(S6, S11)
        S13 = fd.define_scalar(192, dtype=DataType.Int)
        S14 = fd.ops.reciprocal(S3)
        S15 = fd.ops.mul(S13, S14)
        S16 = fd.ops.cast(S12, dtype=DataType.Int)
        S17 = fd.ops.sub(S15, S16)
        T19 = fd.ops.reshape(T1, new_shape=[S0, S17])
        T20 = fd.ops.sum(T19, dims=[1], keepdim=False, dtype=DataType.Null)
        fd.add_output(T20)

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

# Test empty symbolic tensors can be reshaped
# See https://github.com/NVIDIA/Fuser/issues/2362
def test_empty_reshape(nvfuser_direct_test):
    inputs = [torch.randint(0, 10, (0, 1, 2, 3, 4), dtype=torch.int64, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, 1, -1, -1, -1],
            contiguity=[False, None, True, True, True],
            dtype=DataType.Int,
            is_cpu=False,
            stride_order=[4, 3, 2, 1, 0],
        )
        S2 = fd.define_scalar(5, dtype=DataType.Int)
        S3 = fd.define_scalar(0, dtype=DataType.Int)
        T5 = fd.ops.reshape(T0, new_shape=[S2, S3])
        fd.add_output(T5)

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

def test_squeeze(nvfuser_direct_test):
    t0_sizes = [4]
    t1_sizes = [1, 4, 1]
    t2_sizes = [2, 1, 4]
    inputs = [
        torch.randn(*t0_sizes, device="cuda"),
        torch.randn(*t1_sizes, device="cuda"),
        torch.randn(*t2_sizes, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[-1], contiguity=[True])
        t1 = fd.define_tensor(sizes=t1_sizes, strides=[4, 1, 1])
        t2 = fd.define_tensor(sizes=t2_sizes, strides=[4, 4, 1])
        t3 = fd.ops.squeeze(t1, [0, -1])
        t4 = fd.ops.squeeze(t2, [-2])
        t5 = fd.ops.sum(t4, [0])
        t6 = fd.ops.mul(t0, t3)


# See https://github.com/NVIDIA/Fuser/issues/2362
def test_empty_reshape(nvfuser_direct_test):
    inputs = [torch.randint(0, 10, (0, 1, 2, 3, 4), dtype=torch.int64, device="cuda:0")]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, 1, -1, -1, -1],
            contiguity=[False, None, True, True, True],
            dtype=DataType.Int,
            is_cpu=False,
            stride_order=[4, 3, 2, 1, 0],
        )
        S2 = fd.define_scalar(5, dtype=DataType.Int)
        S3 = fd.define_scalar(0, dtype=DataType.Int)
        T5 = fd.ops.reshape(T0, new_shape=[S2, S3])
        fd.add_output(T5)

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

def test_squeeze(nvfuser_direct_test):
    t0_sizes = [4]
    t1_sizes = [1, 4, 1]
    t2_sizes = [2, 1, 4]
    inputs = [
        torch.randn(*t0_sizes, device="cuda"),
        torch.randn(*t1_sizes, device="cuda"),
        torch.randn(*t2_sizes, device="cuda"),
    ]

    def fusion_func(fd: FusionDefinition):
        t0 = fd.define_tensor(shape=[-1], contiguity=[True])
        t1 = fd.define_tensor(sizes=t1_sizes, strides=[4, 1, 1])
        t2 = fd.define_tensor(sizes=t2_sizes, strides=[4, 4, 1])
        t3 = fd.ops.squeeze(t1, [0, -1])
        t4 = fd.ops.squeeze(t2, [-2])
        t5 = fd.ops.sum(t4, [0])
        t6 = fd.ops.mul(t0, t3)
        t7 = fd.ops.mul(t6, t5)
        fd.add_output(t7)

    # Check that squeeze does not print default argument: squeeze_expanded
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[1, -1, 1], contiguity=[None, True, None], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.define_tensor(shape=[-1, 1, -1], contiguity=[True, None, True], dtype=DataType.Float, is_cpu=False)
    tv3 = fd.ops.squeeze(tv1, dims=[0, 2])
    tv6 = fd.ops.mul(tv0, tv3)
    tv4 = fd.ops.squeeze(tv2, dims=[1])
    tv5 = fd.ops.sum(tv4, dims=[0], dtype=DataType.Float)
    tv7 = fd.ops.mul(tv6, tv5)
    fd.add_output(tv7)"""

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

    v1 = torch.sum(inputs[1], [0, -1])
    v2 = torch.sum(inputs[2], [0, 1])
    eager_out = inputs[0] * v1 * v2
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

# Test that expanded dimensions can be reduced properly
# See https://github.com/NVIDIA/Fuser/issues/1678
def test_expanded_reduction(nvfuser_direct_test):
    inputs = [torch.tensor(1.0, device="cuda").as_strided((2, 3), (0, 0))]

    for keepdim in [False, True]:

        def fusion_func(fd: FusionDefinition) -> None:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[None, None],
                dtype=DataType.Float,
                is_cpu=False,
                stride_order=[1, 0],
            )
            T1 = fd.ops.sum(T0, dims=[0], keepdim=keepdim, dtype=DataType.Null)
            fd.add_output(T1)


    tv2 = fd.define_tensor(shape=[-1, 1, -1], contiguity=[True, None, True], dtype=DataType.Float, is_cpu=False)
    tv3 = fd.ops.squeeze(tv1, dims=[0, 2])
    tv6 = fd.ops.mul(tv0, tv3)
    tv4 = fd.ops.squeeze(tv2, dims=[1])
    tv5 = fd.ops.sum(tv4, dims=[0], dtype=DataType.Float)
    tv7 = fd.ops.mul(tv6, tv5)
    fd.add_output(tv7)"""

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

    v1 = torch.sum(inputs[1], [0, -1])
    v2 = torch.sum(inputs[2], [0, 1])
    eager_out = inputs[0] * v1 * v2
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

# Test that expanded dimensions can be reduced properly
# See https://github.com/NVIDIA/Fuser/issues/1678
def test_expanded_reduction(nvfuser_direct_test):
    inputs = [torch.tensor(1.0, device="cuda").as_strided((2, 3), (0, 0))]

    for keepdim in [False, True]:

        def fusion_func(fd: FusionDefinition) -> None:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[None, None],
                dtype=DataType.Float,
                is_cpu=False,
                stride_order=[1, 0],
            )
            T1 = fd.ops.sum(T0, dims=[0], keepdim=keepdim, dtype=DataType.Null)
            fd.add_output(T1)

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

        nvfuser_direct_test.assertEqual(
            nvf_out[0], inputs[0].sum(dim=0, keepdim=keepdim)
        )

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

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

        t0_b = fd.ops.expand(t0, inputs[1].size())
        t2 = fd.ops.mul(t0_b, t1)

        fd.add_output(t2)

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


    for keepdim in [False, True]:

        def fusion_func(fd: FusionDefinition) -> None:
            T0 = fd.define_tensor(
                shape=[-1, -1],
                contiguity=[None, None],
                dtype=DataType.Float,
                is_cpu=False,
                stride_order=[1, 0],
            )
            T1 = fd.ops.sum(T0, dims=[0], keepdim=keepdim, dtype=DataType.Null)
            fd.add_output(T1)

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

        nvfuser_direct_test.assertEqual(
            nvf_out[0], inputs[0].sum(dim=0, keepdim=keepdim)
        )

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

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

        t0_b = fd.ops.expand(t0, inputs[1].size())
        t2 = fd.ops.mul(t0_b, t1)

        fd.add_output(t2)

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

def test_index_select(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
        torch.randn(8, 16, device="cuda"),
        torch.randint(0, 8, (6,), device="cuda").to(dtype=torch.long),
    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            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.index_select(t3, t2, dim)
            fd.add_output(t4)

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


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

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

        t0_b = fd.ops.expand(t0, inputs[1].size())
        t2 = fd.ops.mul(t0_b, t1)

        fd.add_output(t2)

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

def test_index_select(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
        torch.randn(8, 16, device="cuda"),
        torch.randint(0, 8, (6,), device="cuda").to(dtype=torch.long),
    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            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.index_select(t3, t2, dim)
            fd.add_output(t4)

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

        eager_out = torch.index_select(inputs[0] + inputs[1], dim, inputs[2])
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    test_fn(0)
    test_fn(1)

def test_index_select_scalar_indices(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
        torch.tensor(2, device="cuda").to(dtype=torch.long),
    ]

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

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

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


    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            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.index_select(t3, t2, dim)
            fd.add_output(t4)

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

        eager_out = torch.index_select(inputs[0] + inputs[1], dim, inputs[2])
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

    test_fn(0)
    test_fn(1)

def test_index_select_scalar_indices(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
        torch.tensor(2, device="cuda").to(dtype=torch.long),
    ]

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

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

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

    test_fn(0)
    test_fn(1)

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

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            s1 = fd.define_scalar(dtype=DataType.Int)
            t1 = fd.ops.select(t0, s1, dim)
            fd.add_output(t1)

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

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


        torch.randn(8, 16, device="cuda"),
        torch.tensor(2, device="cuda").to(dtype=torch.long),
    ]

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

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

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

    test_fn(0)
    test_fn(1)

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

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            t0 = fd.from_pytorch(inputs[0])
            s1 = fd.define_scalar(dtype=DataType.Int)
            t1 = fd.ops.select(t0, s1, dim)
            fd.add_output(t1)

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

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

    test_fn(0)
    test_fn(1)

def test_take_along_axis(nvfuser_direct_test):
    inputs = [
        torch.randn(8, 16, device="cuda"),
        torch.randn(8, 16, device="cuda"),
        torch.randint(0, 8, (8, 16), device="cuda").to(dtype=torch.long),
    ]

    def test_fn(dim):
        def fusion_func(fd: FusionDefinition):
            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.take_along_axis(t3, t2, dim)
            fd.add_output(t4)

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