tests.python.conftest

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os

def pytest_configure(config):
    """
    Hook called after command line options have been parsed and all plugins
    and initial conftest files have been loaded.
    """
    # Append to NVFUSER_ENABLE environment variable for all tests in this directory
    # Note: id_model is now enabled by default, so we only need extra validation
    existing = os.environ.get("NVFUSER_ENABLE", "")
    new_options = "id_model_extra_validation"

    if existing:
        os.environ["NVFUSER_ENABLE"] = f"{existing},{new_options}"
    else:
        os.environ["NVFUSER_ENABLE"] = new_options

def pytest_make_parametrize_id(config, val, argname):
    return f"{argname}={val}"


import os

def pytest_configure(config):
    """
    Hook called after command line options have been parsed and all plugins
    and initial conftest files have been loaded.
    """
    # Append to NVFUSER_ENABLE environment variable for all tests in this directory
    # Note: id_model is now enabled by default, so we only need extra validation
    existing = os.environ.get("NVFUSER_ENABLE", "")
    new_options = "id_model_extra_validation"

    if existing:
        os.environ["NVFUSER_ENABLE"] = f"{existing},{new_options}"
    else:
        os.environ["NVFUSER_ENABLE"] = new_options

def pytest_make_parametrize_id(config, val, argname):
    return f"{argname}={val}"


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import math
from contextlib import contextmanager
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from thunder.dynamo import thunderfx

# Sizes used in Llama 4 Maverick
@dataclass
class Config:
    hidden_size: int = 5120
    intermediate_size: int = 8192
    num_routed_experts: int = 128
    num_shared_experts: int = 1

class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )

def _group_sizes_from_offsets(offsets: torch.Tensor) -> list[int]:
    group_sizes = []
    prev = 0
    for offset in offsets:
        group_sizes.append(offset - prev)
        prev = offset
    return group_sizes

# Required otherwise, there is a graph-break.
_grouped_mm = torch.compiler.allow_in_graph(torch._grouped_mm)

# This function should be replaced with torch._grouped_mm.  However,
# torch._grouped_mm is yet to be usable because it requires offsets being
# multiples of 16.
def grouped_mm(a: torch.Tensor, b: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    if torch.compiler.is_compiling():
        return _grouped_mm(a, b, offsets)

    group_sizes = _group_sizes_from_offsets(offsets)
    group_outs = []
    for group_a, group_b in zip(a.split(group_sizes), b.unbind()):
        group_outs.append(group_a @ group_b)
    return torch.cat(group_outs)

class GroupedLinear(nn.Module):
    def __init__(self, groups: int, in_features: int, out_features: int):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))
        # Initialize the weight in the same way as nn.Linear
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return grouped_mm(hidden_states, self.weight, offsets)

class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.down_proj = GroupedLinear(groups, intermediate_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_routed_experts, bias=False)
        self.shared_experts = SwiGLU(
            config.hidden_size, config.intermediate_size * config.num_shared_experts
        )
        self.routed_experts = GroupedSwiGLU(
            config.num_routed_experts, config.hidden_size, config.intermediate_size
        )

    def run_routed_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))  # [s, h]

        router_logits = self.gate(hidden_states)  # [s, n]
        topk_weight, topk_ids = router_logits.topk(1)  # [s, 1]
        router_scores = topk_weight.sigmoid()  # [s, 1]
        hidden_states = hidden_states * router_scores  # [s, h]

        counts = torch.zeros(
            topk_ids.size(0),
            self.config.num_routed_experts,
            device=topk_ids.device,
            dtype=torch.int32,
        )  # [s, n]
        counts = counts.scatter(1, topk_ids, 1)  # [s, n]
        tokens_per_expert = counts.sum(0)  # [n]

        token_ids_sorted_by_expert_id = topk_ids.view(-1).argsort()  # [s]
        tokens_sorted_by_expert_id = hidden_states[
            token_ids_sorted_by_expert_id
        ]  # [s, h]

        # Without `torch.int32`, we see `RuntimeError: Offsets tensor must be integer (int32) tensor, but got torch.int64.`
        # from PyTorch when calling _grouped_mm.
        offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)  # [n]
        outs_sorted_by_expert_id = self.routed_experts(
            tokens_sorted_by_expert_id, offsets
        )  # [s, h]

        token_ids_sorted_by_expert_inverse_id = torch.argsort(
            token_ids_sorted_by_expert_id
        )
        outs_sorted_by_token_id = outs_sorted_by_expert_id[
            token_ids_sorted_by_expert_inverse_id
        ]

        return outs_sorted_by_token_id

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.shared_experts(hidden_states) + self.run_routed_experts(
            hidden_states
        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore
    torch.set_default_dtype(prev_dtype)
    torch.set_default_device(prev_device)

def test_llama4_moe_thunderfx():
    config = Config()

    # This is much faster than creating the module with CPU float parameters
    # and then doing `.to("cuda").to(torch.bfloat16)`.
    with default_tensor_type(dtype=torch.bfloat16, device="cuda"):
        model = Llama4MoE(config)

    # Without this, `thunderfx` falls back to `inductor` for `_grouped_mm`
    # as it doesn't have a grad-rule for the same.
    model.requires_grad_(False)

    batch_size, seq_len = 1, 2048
    inp = torch.randn(
        batch_size, seq_len, config.hidden_size, dtype=torch.bfloat16, device="cuda"
    )
    expected = model(inp)

    assert expected.size() == (batch_size, seq_len, config.hidden_size)
    assert expected.dtype == torch.bfloat16
    assert expected.is_cuda

    tmodel = thunderfx(model, nv_enable_linear=True, nv_enable_scatter=True)

    with torch.no_grad():
        actual = tmodel(inp)

    assert len(tmodel._backend.subgraph_infos) == 1
    assert len(tmodel._backend.subgraph_infos[0].split_reasons) == 0
    # Uncomment to view thunder traces
    # print(tmodel.last_traces)

    torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)


Config

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import math
from contextlib import contextmanager
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from thunder.dynamo import thunderfx

# Sizes used in Llama 4 Maverick
@dataclass
class Config:
    hidden_size: int = 5120
    intermediate_size: int = 8192
    num_routed_experts: int = 128
    num_shared_experts: int = 1

class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )

def _group_sizes_from_offsets(offsets: torch.Tensor) -> list[int]:
    group_sizes = []
    prev = 0
    for offset in offsets:
        group_sizes.append(offset - prev)


SwiGLU

import math
from contextlib import contextmanager
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from thunder.dynamo import thunderfx

# Sizes used in Llama 4 Maverick
@dataclass
class Config:
    hidden_size: int = 5120
    intermediate_size: int = 8192
    num_routed_experts: int = 128
    num_shared_experts: int = 1

class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )

def _group_sizes_from_offsets(offsets: torch.Tensor) -> list[int]:
    group_sizes = []
    prev = 0
    for offset in offsets:
        group_sizes.append(offset - prev)
        prev = offset
    return group_sizes

# Required otherwise, there is a graph-break.
_grouped_mm = torch.compiler.allow_in_graph(torch._grouped_mm)

# This function should be replaced with torch._grouped_mm.  However,
# torch._grouped_mm is yet to be usable because it requires offsets being
# multiples of 16.
def grouped_mm(a: torch.Tensor, b: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    if torch.compiler.is_compiling():


class Config:
    hidden_size: int = 5120
    intermediate_size: int = 8192
    num_routed_experts: int = 128
    num_shared_experts: int = 1

class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )

def _group_sizes_from_offsets(offsets: torch.Tensor) -> list[int]:
    group_sizes = []
    prev = 0
    for offset in offsets:
        group_sizes.append(offset - prev)
        prev = offset
    return group_sizes

# Required otherwise, there is a graph-break.
_grouped_mm = torch.compiler.allow_in_graph(torch._grouped_mm)

# This function should be replaced with torch._grouped_mm.  However,
# torch._grouped_mm is yet to be usable because it requires offsets being
# multiples of 16.
def grouped_mm(a: torch.Tensor, b: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    if torch.compiler.is_compiling():
        return _grouped_mm(a, b, offsets)

    group_sizes = _group_sizes_from_offsets(offsets)
    group_outs = []
    for group_a, group_b in zip(a.split(group_sizes), b.unbind()):
        group_outs.append(group_a @ group_b)
    return torch.cat(group_outs)


            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )

def _group_sizes_from_offsets(offsets: torch.Tensor) -> list[int]:
    group_sizes = []
    prev = 0
    for offset in offsets:
        group_sizes.append(offset - prev)
        prev = offset
    return group_sizes

# Required otherwise, there is a graph-break.
_grouped_mm = torch.compiler.allow_in_graph(torch._grouped_mm)

# This function should be replaced with torch._grouped_mm.  However,
# torch._grouped_mm is yet to be usable because it requires offsets being
# multiples of 16.
def grouped_mm(a: torch.Tensor, b: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    if torch.compiler.is_compiling():
        return _grouped_mm(a, b, offsets)

    group_sizes = _group_sizes_from_offsets(offsets)
    group_outs = []
    for group_a, group_b in zip(a.split(group_sizes), b.unbind()):
        group_outs.append(group_a @ group_b)
    return torch.cat(group_outs)

class GroupedLinear(nn.Module):
    def __init__(self, groups: int, in_features: int, out_features: int):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))
        # Initialize the weight in the same way as nn.Linear
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return grouped_mm(hidden_states, self.weight, offsets)

class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)


GroupedLinear

# Required otherwise, there is a graph-break.
_grouped_mm = torch.compiler.allow_in_graph(torch._grouped_mm)

# This function should be replaced with torch._grouped_mm.  However,
# torch._grouped_mm is yet to be usable because it requires offsets being
# multiples of 16.
def grouped_mm(a: torch.Tensor, b: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    if torch.compiler.is_compiling():
        return _grouped_mm(a, b, offsets)

    group_sizes = _group_sizes_from_offsets(offsets)
    group_outs = []
    for group_a, group_b in zip(a.split(group_sizes), b.unbind()):
        group_outs.append(group_a @ group_b)
    return torch.cat(group_outs)

class GroupedLinear(nn.Module):
    def __init__(self, groups: int, in_features: int, out_features: int):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))
        # Initialize the weight in the same way as nn.Linear
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return grouped_mm(hidden_states, self.weight, offsets)

class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.down_proj = GroupedLinear(groups, intermediate_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):


GroupedSwiGLU

    group_sizes = _group_sizes_from_offsets(offsets)
    group_outs = []
    for group_a, group_b in zip(a.split(group_sizes), b.unbind()):
        group_outs.append(group_a @ group_b)
    return torch.cat(group_outs)

class GroupedLinear(nn.Module):
    def __init__(self, groups: int, in_features: int, out_features: int):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))
        # Initialize the weight in the same way as nn.Linear
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return grouped_mm(hidden_states, self.weight, offsets)

class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.down_proj = GroupedLinear(groups, intermediate_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_routed_experts, bias=False)
        self.shared_experts = SwiGLU(
            config.hidden_size, config.intermediate_size * config.num_shared_experts
        )
        self.routed_experts = GroupedSwiGLU(
            config.num_routed_experts, config.hidden_size, config.intermediate_size
        )

    def run_routed_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))  # [s, h]

        router_logits = self.gate(hidden_states)  # [s, n]
        topk_weight, topk_ids = router_logits.topk(1)  # [s, 1]


        return grouped_mm(hidden_states, self.weight, offsets)

class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.down_proj = GroupedLinear(groups, intermediate_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_routed_experts, bias=False)
        self.shared_experts = SwiGLU(
            config.hidden_size, config.intermediate_size * config.num_shared_experts
        )
        self.routed_experts = GroupedSwiGLU(
            config.num_routed_experts, config.hidden_size, config.intermediate_size
        )

    def run_routed_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))  # [s, h]

        router_logits = self.gate(hidden_states)  # [s, n]
        topk_weight, topk_ids = router_logits.topk(1)  # [s, 1]
        router_scores = topk_weight.sigmoid()  # [s, 1]
        hidden_states = hidden_states * router_scores  # [s, h]

        counts = torch.zeros(
            topk_ids.size(0),
            self.config.num_routed_experts,
            device=topk_ids.device,
            dtype=torch.int32,
        )  # [s, n]
        counts = counts.scatter(1, topk_ids, 1)  # [s, n]
        tokens_per_expert = counts.sum(0)  # [n]

        token_ids_sorted_by_expert_id = topk_ids.view(-1).argsort()  # [s]
        tokens_sorted_by_expert_id = hidden_states[
            token_ids_sorted_by_expert_id
        ]  # [s, h]

        # Without `torch.int32`, we see `RuntimeError: Offsets tensor must be integer (int32) tensor, but got torch.int64.`
        # from PyTorch when calling _grouped_mm.
        offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)  # [n]
        outs_sorted_by_expert_id = self.routed_experts(
            tokens_sorted_by_expert_id, offsets
        )  # [s, h]

        token_ids_sorted_by_expert_inverse_id = torch.argsort(
            token_ids_sorted_by_expert_id
        )
        outs_sorted_by_token_id = outs_sorted_by_expert_id[
            token_ids_sorted_by_expert_inverse_id
        ]

        return outs_sorted_by_token_id

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.shared_experts(hidden_states) + self.run_routed_experts(
            hidden_states
        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore
    torch.set_default_dtype(prev_dtype)
    torch.set_default_device(prev_device)

def test_llama4_moe_thunderfx():


        outs_sorted_by_expert_id = self.routed_experts(
            tokens_sorted_by_expert_id, offsets
        )  # [s, h]

        token_ids_sorted_by_expert_inverse_id = torch.argsort(
            token_ids_sorted_by_expert_id
        )
        outs_sorted_by_token_id = outs_sorted_by_expert_id[
            token_ids_sorted_by_expert_inverse_id
        ]

        return outs_sorted_by_token_id

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.shared_experts(hidden_states) + self.run_routed_experts(
            hidden_states
        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore
    torch.set_default_dtype(prev_dtype)
    torch.set_default_device(prev_device)

def test_llama4_moe_thunderfx():
    config = Config()

    # This is much faster than creating the module with CPU float parameters
    # and then doing `.to("cuda").to(torch.bfloat16)`.
    with default_tensor_type(dtype=torch.bfloat16, device="cuda"):
        model = Llama4MoE(config)

    # Without this, `thunderfx` falls back to `inductor` for `_grouped_mm`
    # as it doesn't have a grad-rule for the same.
    model.requires_grad_(False)

    batch_size, seq_len = 1, 2048
    inp = torch.randn(
        batch_size, seq_len, config.hidden_size, dtype=torch.bfloat16, device="cuda"
    )
    expected = model(inp)


        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore
    torch.set_default_dtype(prev_dtype)
    torch.set_default_device(prev_device)

def test_llama4_moe_thunderfx():
    config = Config()

    # This is much faster than creating the module with CPU float parameters
    # and then doing `.to("cuda").to(torch.bfloat16)`.
    with default_tensor_type(dtype=torch.bfloat16, device="cuda"):
        model = Llama4MoE(config)

    # Without this, `thunderfx` falls back to `inductor` for `_grouped_mm`
    # as it doesn't have a grad-rule for the same.
    model.requires_grad_(False)

    batch_size, seq_len = 1, 2048
    inp = torch.randn(
        batch_size, seq_len, config.hidden_size, dtype=torch.bfloat16, device="cuda"
    )
    expected = model(inp)

    assert expected.size() == (batch_size, seq_len, config.hidden_size)
    assert expected.dtype == torch.bfloat16
    assert expected.is_cuda

    tmodel = thunderfx(model, nv_enable_linear=True, nv_enable_scatter=True)

    with torch.no_grad():
        actual = tmodel(inp)

    assert len(tmodel._backend.subgraph_infos) == 1
    assert len(tmodel._backend.subgraph_infos[0].split_reasons) == 0
    # Uncomment to view thunder traces
    # print(tmodel.last_traces)

    torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)


class GroupedSwiGLU(nn.Module):
    def __init__(self, groups: int, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.up_proj = GroupedLinear(groups, hidden_size, intermediate_size)
        self.down_proj = GroupedLinear(groups, intermediate_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, offsets: torch.Tensor
    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_routed_experts, bias=False)
        self.shared_experts = SwiGLU(
            config.hidden_size, config.intermediate_size * config.num_shared_experts
        )
        self.routed_experts = GroupedSwiGLU(
            config.num_routed_experts, config.hidden_size, config.intermediate_size
        )

    def run_routed_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))  # [s, h]

        router_logits = self.gate(hidden_states)  # [s, n]
        topk_weight, topk_ids = router_logits.topk(1)  # [s, 1]
        router_scores = topk_weight.sigmoid()  # [s, 1]
        hidden_states = hidden_states * router_scores  # [s, h]

        counts = torch.zeros(
            topk_ids.size(0),
            self.config.num_routed_experts,
            device=topk_ids.device,
            dtype=torch.int32,
        )  # [s, n]
        counts = counts.scatter(1, topk_ids, 1)  # [s, n]
        tokens_per_expert = counts.sum(0)  # [n]

        token_ids_sorted_by_expert_id = topk_ids.view(-1).argsort()  # [s]


forward

        tokens_sorted_by_expert_id = hidden_states[
            token_ids_sorted_by_expert_id
        ]  # [s, h]

        # Without `torch.int32`, we see `RuntimeError: Offsets tensor must be integer (int32) tensor, but got torch.int64.`
        # from PyTorch when calling _grouped_mm.
        offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)  # [n]
        outs_sorted_by_expert_id = self.routed_experts(
            tokens_sorted_by_expert_id, offsets
        )  # [s, h]

        token_ids_sorted_by_expert_inverse_id = torch.argsort(
            token_ids_sorted_by_expert_id
        )
        outs_sorted_by_token_id = outs_sorted_by_expert_id[
            token_ids_sorted_by_expert_inverse_id
        ]

        return outs_sorted_by_token_id

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.shared_experts(hidden_states) + self.run_routed_experts(
            hidden_states
        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore
    torch.set_default_dtype(prev_dtype)
    torch.set_default_device(prev_device)

def test_llama4_moe_thunderfx():


    ) -> torch.Tensor:
        return self.down_proj(
            F.silu(self.gate_proj(hidden_states, offsets))
            * self.up_proj(hidden_states, offsets),
            offsets,
        )

class Llama4MoE(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_routed_experts, bias=False)
        self.shared_experts = SwiGLU(
            config.hidden_size, config.intermediate_size * config.num_shared_experts
        )
        self.routed_experts = GroupedSwiGLU(
            config.num_routed_experts, config.hidden_size, config.intermediate_size
        )

    def run_routed_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))  # [s, h]

        router_logits = self.gate(hidden_states)  # [s, n]
        topk_weight, topk_ids = router_logits.topk(1)  # [s, 1]
        router_scores = topk_weight.sigmoid()  # [s, 1]
        hidden_states = hidden_states * router_scores  # [s, h]

        counts = torch.zeros(
            topk_ids.size(0),
            self.config.num_routed_experts,
            device=topk_ids.device,
            dtype=torch.int32,
        )  # [s, n]
        counts = counts.scatter(1, topk_ids, 1)  # [s, n]
        tokens_per_expert = counts.sum(0)  # [n]

        token_ids_sorted_by_expert_id = topk_ids.view(-1).argsort()  # [s]
        tokens_sorted_by_expert_id = hidden_states[
            token_ids_sorted_by_expert_id
        ]  # [s, h]

        # Without `torch.int32`, we see `RuntimeError: Offsets tensor must be integer (int32) tensor, but got torch.int64.`
        # from PyTorch when calling _grouped_mm.
        offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)  # [n]
        outs_sorted_by_expert_id = self.routed_experts(
            tokens_sorted_by_expert_id, offsets
        )  # [s, h]

        token_ids_sorted_by_expert_inverse_id = torch.argsort(
            token_ids_sorted_by_expert_id
        )
        outs_sorted_by_token_id = outs_sorted_by_expert_id[
            token_ids_sorted_by_expert_inverse_id
        ]

        return outs_sorted_by_token_id

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.shared_experts(hidden_states) + self.run_routed_experts(
            hidden_states
        )

@contextmanager
def default_tensor_type(dtype=torch.float32, device="cpu"):
    # Save
    prev_dtype = torch.get_default_dtype()
    prev_device = torch.get_default_device()

    # Set
    torch.set_default_dtype(dtype)
    torch.set_default_device(device)

    yield

    # Restore


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

import torch

from opinfo_core import OpInfo
from opinfo_utils import ArgumentType, is_tensor

import sys

if "nvfuser" in sys.modules:
    from nvfuser import FusionDefinition
    from nvfuser.pytorch_utils import (
        python_scalar_to_nvfuser_dtype,
        torch_dtype_to_nvfuser_dtype,
    )
else:
    from nvfuser_direct import FusionDefinition
    from nvfuser_direct.pytorch_utils import (
        python_scalar_to_nvfuser_dtype,
        torch_dtype_to_nvfuser_dtype,
    )

def parse_inputs_fusion_definition(fd: FusionDefinition, opinfo: OpInfo, *args):
    if len(args) == 0:
        return []

    nvf_args = []

    symbolic_parameter_list = (
        opinfo.symbolic_parameter_list
        if opinfo.symbolic_parameter_list is not None
        else [ArgumentType.Symbolic] * len(args)
    )

    num_symbolic_parameters = len(symbolic_parameter_list)
    assert num_symbolic_parameters >= len(
        args
    ), f"{num_symbolic_parameters} vs {len(args)}"

    for arg_type, a in zip(symbolic_parameter_list, args):
        if arg_type == ArgumentType.Symbolic:
            if isinstance(a, torch.Tensor):
                nvf_args.append(fd.from_pytorch(a))
            elif isinstance(a, list) and all(map(is_tensor, a)):
                nvf_args.append([fd.from_pytorch(inner_a) for inner_a in a])
            elif isinstance(a, list) or isinstance(a, tuple):
                nvf_args.append(fd.define_vector(a))
            else:
                # For symbolic scalars, we do not define with constant value.
                # Otherwise, it becomes a constant and is not a fusion input.
                nvf_args.append(fd.define_scalar(python_scalar_to_nvfuser_dtype(a)))
        elif arg_type == ArgumentType.ConstantScalar:
            assert not isinstance(a, torch.Tensor)
            nvf_args.append(fd.define_scalar(a))
        elif isinstance(a, torch.dtype):
            nvf_args.append(torch_dtype_to_nvfuser_dtype(a))
        else:
            assert not isinstance(a, torch.Tensor)
            assert arg_type == ArgumentType.Constant
            nvf_args.append(a)
    return nvf_args

# This function will purposely not generate a functional FusionDefintion as
# it lacks defining an output.  It is only meant to test the error checking
# of an operation.
def api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)

def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)

def tensor_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    out = opinfo.op(fd)(nvf_inputs[0], **kwargs)

def vector_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    v0 = nvf_inputs[0].shape()
    out = opinfo.op(fd)(v0, **kwargs)


from opinfo_core import OpInfo
from opinfo_utils import ArgumentType, is_tensor

import sys

if "nvfuser" in sys.modules:
    from nvfuser import FusionDefinition
    from nvfuser.pytorch_utils import (
        python_scalar_to_nvfuser_dtype,
        torch_dtype_to_nvfuser_dtype,
    )
else:
    from nvfuser_direct import FusionDefinition
    from nvfuser_direct.pytorch_utils import (
        python_scalar_to_nvfuser_dtype,
        torch_dtype_to_nvfuser_dtype,
    )

def parse_inputs_fusion_definition(fd: FusionDefinition, opinfo: OpInfo, *args):
    if len(args) == 0:
        return []

    nvf_args = []

    symbolic_parameter_list = (
        opinfo.symbolic_parameter_list
        if opinfo.symbolic_parameter_list is not None
        else [ArgumentType.Symbolic] * len(args)
    )

    num_symbolic_parameters = len(symbolic_parameter_list)
    assert num_symbolic_parameters >= len(
        args
    ), f"{num_symbolic_parameters} vs {len(args)}"

    for arg_type, a in zip(symbolic_parameter_list, args):
        if arg_type == ArgumentType.Symbolic:
            if isinstance(a, torch.Tensor):
                nvf_args.append(fd.from_pytorch(a))
            elif isinstance(a, list) and all(map(is_tensor, a)):
                nvf_args.append([fd.from_pytorch(inner_a) for inner_a in a])
            elif isinstance(a, list) or isinstance(a, tuple):
                nvf_args.append(fd.define_vector(a))
            else:
                # For symbolic scalars, we do not define with constant value.
                # Otherwise, it becomes a constant and is not a fusion input.
                nvf_args.append(fd.define_scalar(python_scalar_to_nvfuser_dtype(a)))
        elif arg_type == ArgumentType.ConstantScalar:
            assert not isinstance(a, torch.Tensor)
            nvf_args.append(fd.define_scalar(a))
        elif isinstance(a, torch.dtype):
            nvf_args.append(torch_dtype_to_nvfuser_dtype(a))
        else:
            assert not isinstance(a, torch.Tensor)
            assert arg_type == ArgumentType.Constant
            nvf_args.append(a)
    return nvf_args

# This function will purposely not generate a functional FusionDefintion as
# it lacks defining an output.  It is only meant to test the error checking
# of an operation.
def api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)

def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)


                nvf_args.append(fd.define_vector(a))
            else:
                # For symbolic scalars, we do not define with constant value.
                # Otherwise, it becomes a constant and is not a fusion input.
                nvf_args.append(fd.define_scalar(python_scalar_to_nvfuser_dtype(a)))
        elif arg_type == ArgumentType.ConstantScalar:
            assert not isinstance(a, torch.Tensor)
            nvf_args.append(fd.define_scalar(a))
        elif isinstance(a, torch.dtype):
            nvf_args.append(torch_dtype_to_nvfuser_dtype(a))
        else:
            assert not isinstance(a, torch.Tensor)
            assert arg_type == ArgumentType.Constant
            nvf_args.append(a)
    return nvf_args

# This function will purposely not generate a functional FusionDefintion as
# it lacks defining an output.  It is only meant to test the error checking
# of an operation.
def api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)

def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)


        elif arg_type == ArgumentType.ConstantScalar:
            assert not isinstance(a, torch.Tensor)
            nvf_args.append(fd.define_scalar(a))
        elif isinstance(a, torch.dtype):
            nvf_args.append(torch_dtype_to_nvfuser_dtype(a))
        else:
            assert not isinstance(a, torch.Tensor)
            assert arg_type == ArgumentType.Constant
            nvf_args.append(a)
    return nvf_args

# This function will purposely not generate a functional FusionDefintion as
# it lacks defining an output.  It is only meant to test the error checking
# of an operation.
def api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)

def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)

def tensor_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    out = opinfo.op(fd)(nvf_inputs[0], **kwargs)

def vector_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    v0 = nvf_inputs[0].shape()
    out = opinfo.op(fd)(v0, **kwargs)


# This function will purposely not generate a functional FusionDefintion as
# it lacks defining an output.  It is only meant to test the error checking
# of an operation.
def api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)

def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)

def tensor_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    out = opinfo.op(fd)(nvf_inputs[0], **kwargs)

def vector_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    v0 = nvf_inputs[0].shape()
    out = opinfo.op(fd)(v0, **kwargs)


def default_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    result = opinfo.op(fd)(*nvf_inputs, **kwargs)
    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)

def tensor_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    out = opinfo.op(fd)(nvf_inputs[0], **kwargs)

def vector_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    v0 = nvf_inputs[0].shape()
    out = opinfo.op(fd)(v0, **kwargs)


    if isinstance(result, tuple):
        for a in result:
            if a is not None:
                fd.add_output(a)
    else:
        fd.add_output(result)

def tensor_input_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    this_inputs = opinfo.op(fd)(**kwargs)
    t1 = fd.ops.add(nvf_inputs[0], this_inputs)
    fd.add_output(t1)

def tensor_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    out = opinfo.op(fd)(nvf_inputs[0], **kwargs)

def vector_api_test_fd_fn(fd: FusionDefinition, opinfo: OpInfo, *args, **kwargs):
    nvf_inputs = parse_inputs_fusion_definition(fd, opinfo, *args)
    v0 = nvf_inputs[0].shape()
    out = opinfo.op(fd)(v0, **kwargs)


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

from opinfo_utils import (
    all_dtypes_except_reduced,
    ArgumentType,
    torch_to_python_dtype_map,
)
from typing import Callable, Optional
import torch
from enum import Enum
from dataclasses import dataclass, field

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3

@dataclass
class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""

    op: Callable

    name: str

    # Set of valid inputs for this operation
    domain: Domain = field(default_factory=lambda: Domain(None, None))

    # Set of valid dtypes for this operation
    dtypes: tuple = all_dtypes_except_reduced

    # Generates valid inputs
    sample_input_generator: Callable = None

    # Generates error inputs
    error_input_generator: Callable = None

    # Function of FusionDefintion operations for valid inputs
    fd_correctness_fn: Callable = None

    # Function of FusionDefintion operations for error inputs
    fd_error_input_fn: Callable = None

    # Reference function for operation
    reference: Callable = None

    # Designate which framework defines the reference
    reference_type: ReferenceType = ReferenceType.Pytorch

    # Nvfuser requires reduction axes to be constant values.
    # symbolic_parameter_list specifies whether an operation's parameters are symbolic.
    # All keyword arguments are considered constant.
    # If symbolic_parameter_list is None, then we assume all parameters to be symbolic.
    symbolic_parameter_list: Optional[list[ArgumentType]] = None

    # Enable check_cpp_translation test
    # Tests that translation from CPP Fusion back to Python FusionDefinition is correct.
    is_clonable: bool = False

    # Test operation using direct bindings python API
    supports_direct_bindings: bool = True


ReferenceType

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

from opinfo_utils import (
    all_dtypes_except_reduced,
    ArgumentType,
    torch_to_python_dtype_map,
)
from typing import Callable, Optional
import torch
from enum import Enum
from dataclasses import dataclass, field

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3

@dataclass
class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",


ErrorSample

from opinfo_utils import (
    all_dtypes_except_reduced,
    ArgumentType,
    torch_to_python_dtype_map,
)
from typing import Callable, Optional
import torch
from enum import Enum
from dataclasses import dataclass, field

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3

@dataclass
class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs


Domain

import torch
from enum import Enum
from dataclasses import dataclass, field

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3

@dataclass
class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE


SampleInput

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3

@dataclass
class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""

    op: Callable

    name: str

    # Set of valid inputs for this operation
    domain: Domain = field(default_factory=lambda: Domain(None, None))

    # Set of valid dtypes for this operation
    dtypes: tuple = all_dtypes_except_reduced

    # Generates valid inputs
    sample_input_generator: Callable = None

    # Generates error inputs


OpInfo

        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""

    op: Callable

    name: str

    # Set of valid inputs for this operation
    domain: Domain = field(default_factory=lambda: Domain(None, None))

    # Set of valid dtypes for this operation
    dtypes: tuple = all_dtypes_except_reduced

    # Generates valid inputs
    sample_input_generator: Callable = None

    # Generates error inputs
    error_input_generator: Callable = None

    # Function of FusionDefintion operations for valid inputs
    fd_correctness_fn: Callable = None

    # Function of FusionDefintion operations for error inputs
    fd_error_input_fn: Callable = None

    # Reference function for operation
    reference: Callable = None

    # Designate which framework defines the reference
    reference_type: ReferenceType = ReferenceType.Pytorch

    # Nvfuser requires reduction axes to be constant values.
    # symbolic_parameter_list specifies whether an operation's parameters are symbolic.
    # All keyword arguments are considered constant.
    # If symbolic_parameter_list is None, then we assume all parameters to be symbolic.
    symbolic_parameter_list: Optional[list[ArgumentType]] = None

    # Enable check_cpp_translation test
    # Tests that translation from CPP Fusion back to Python FusionDefinition is correct.
    is_clonable: bool = False

    # Test operation using direct bindings python API
    supports_direct_bindings: bool = True


class ErrorSample:
    kwargs: dict
    ex_str: str
    ex_type: Exception = RuntimeError

@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.


@dataclass
class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())


jax

class Domain:
    low: int
    high: int

class SampleInput:
    """Represents sample inputs to a function."""

    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""


python

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""

    op: Callable

    name: str

    # Set of valid inputs for this operation
    domain: Domain = field(default_factory=lambda: Domain(None, None))

    # Set of valid dtypes for this operation
    dtypes: tuple = all_dtypes_except_reduced

    # Generates valid inputs
    sample_input_generator: Callable = None

    # Generates error inputs


    __slots__ = [
        "args",
        "kwargs",
    ]

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return f"[SampleInput args={self.args} kwargs={self.kwargs}]"

    def jax(self):
        from opinfo_utils import JAX_AVAILABLE

        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())


        assert JAX_AVAILABLE

        import jax.numpy as jnp
        from opinfo_utils import torch_to_jax_dtype_map

        def to_jax(t):
            if isinstance(t, torch.Tensor):
                return jnp.array(t.cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_jax_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_jax, self.args)
        return SampleInput(*args, *self.kwargs.values())

    def python(self):
        # Flatten Pytorch Tensors into Python Lists
        def to_python(t):
            if isinstance(t, torch.Tensor):
                return list(t.flatten().cpu().numpy())
            if isinstance(t, torch.dtype):
                return torch_to_python_dtype_map[t]
            return t

        # Note: We assume arguments have flat hierarchy.
        # TODO Add support for kwargs
        args = map(to_python, self.args)
        return SampleInput(*args, *self.kwargs.values())

@dataclass
class OpInfo:
    """Operator information and helper functions for acquiring it."""

    op: Callable

    name: str

    # Set of valid inputs for this operation
    domain: Domain = field(default_factory=lambda: Domain(None, None))

    # Set of valid dtypes for this operation
    dtypes: tuple = all_dtypes_except_reduced


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

import itertools
from functools import partial, wraps

import math
import torch
from torch.testing import make_tensor
import random
from numbers import Number

from opinfo_core import OpInfo, SampleInput, ErrorSample, Domain
from opinfo_utils import (
    make_number,
    find_nonmatching_dtype,
    is_floating_dtype,
    float_complex_dtypes,
    complex_dtypes,
)
import sys

if "nvfuser" in sys.modules:
    from nvfuser import DataType
    from nvfuser.pytorch_utils import torch_dtype_to_nvfuser_dtype
else:
    from nvfuser_direct import DataType
    from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

MINIMUM_SYMBOLIC_SIZE = -1
INT64_MAX = 2**63 - 1
MAX_TENSOR_DIMS = 8

# Determine if a number is with desired Domain [low, high)
# The domain is half-open. The lower limit is inclusive while the upper limit is exclusive.
def is_within_domain(domain: Domain, a: Number, exclude_zero: bool = False):
    # comparison operators are not defined for complex numbers
    if isinstance(a, complex):
        return True

    if domain.low is not None and domain.low > a:
        return False

    if domain.high is not None and a >= domain.high:
        return False

    if exclude_zero and a == 0:
        return False

    return True

def _extremal_values(dtype: torch.dtype):
    _float_vals = (float("inf"), float("-inf"), float("nan"))
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )
    _int16_vals = (-32768, 32767)
    _int32_vals = (-2147483648, 2147483647)
    _int64_vals = (-9223372036854775808, 9223372036854775807)

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype is torch.int16:
        return _int16_vals
    elif dtype is torch.int32:
        return _int32_vals
    elif dtype is torch.int64:
        return _int64_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _large_values(dtype: torch.dtype):
    _int_vals = (-1113, 1113, -10701, 10701)
    _float16_vals = (-501, 501, -1001.2, 1001.2, -13437.7, 13437.7)
    _float_vals = _float16_vals + (-4988429.2, 4988429.2, -1e20, 1e20)
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype == torch.float16:
        return _float16_vals
    elif dtype in (torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        print(_complex_vals)
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _small_values(dtype: torch.dtype):
    eps = 1e-5
    _int_vals = (0, -1, 1, -55, 55, -127, 127, -128)
    _float_vals = (
        0.0,
        -0.0,
        -1e-3,
        1e-3,
        -0.25,
        0.25,
        -1.0,
        1.0,
        -math.e / 2.0,
        math.e / 2.0,
        -math.e + eps,
        math.e - eps,
        -math.e,
        math.e,
        -math.e - eps,
        math.e + eps,
    )
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def broadcast_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast(operand, sizes)
    # add new dimensions to left-hand-side of tensor
    # dims = tuple(range(len(sizes), len(sizes) + np.ndim(operand)))
    # return broadcast_in_dim(operand, tuple(sizes) + np.shape(operand), dims)

    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    fewer_original_axes = (
        ([2, 3], [True, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    greater_original_axes = (
        ([2, 3], [True, False, False, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    error_cases = [
        fewer_original_axes,
        greater_original_axes,
    ]
    for es in error_cases:
        ex_case, ex_type, ex_str = es
        input_shape, bcast_dims = ex_case
        input_tensor = make_arg(input_shape)
        yield SampleInput(input_tensor, bcast_dims), ex_type, ex_str

def broadcast_in_dim_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # The first 5 test cases below are taken from JAX's broadcast_in_dim tests
    #   https://github.com/google/jax/blob/main/tests/lax_test.py#L1171
    # input shape, output shape, bcast_dims
    cases = (
        ([2], [2, 2], [0]),
        ([2], [2, 2], [1]),
        ([2], [2, 3], [0]),
        ([], [2, 3], []),
        ([1], [2, 3], [1]),
        ((4, 6, 3, 1), (5, 4, 7, 6, 3, 6, 6), (1, 3, 4, 5)),
    )

    for input_shape, output_shape, bcast_dims in cases:
        input_tensor = make_arg(input_shape)
        if op.name == "broadcast_in_dim_symbolic":
            bcast_shaped_tensor = make_arg(output_shape)
            yield SampleInput(input_tensor, bcast_shaped_tensor, bcast_dims)
        else:
            yield SampleInput(input_tensor, output_shape, bcast_dims)

def broadcast_in_dim_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # 1. Every dimension in the input tensor must be used in broadcast_dimensions.
    missing_axis_in_bcast_dims = (
        ([2, 2], [2, 2, 3], [0]),
        RuntimeError,
        "The broadcast dimensions should match the input dimensions.",
    )

    # 2. New shape has weakly more dimentions than the original tensor.
    fewer_dims_in_output_shape = (
        ([2, 2], [2], [0]),
        RuntimeError,
        "The new shape is expected to be greater-then-or-equal to the input",
    )

    # 3. Each broadcast dimension is within the new shape.
    out_of_bounds_broadcast_dimensions = (
        ([2, 2], [2, 2], [0, 2]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # 4. The original tensor is not broadcastable to desired shape.
    # tensor.shape[idx] == 1 or tensor.shape[idx] == output_shape[new_idx]
    #
    # Jax Exception:
    # TypeError: broadcast_in_dim operand dimension sizes must either be 1,
    # or be equal to their corresponding dimensions in the target broadcast shape;
    # got operand of shape (2, 3), target broadcast shape (2, 3, 4), broadcast_dimensions (0, 2)
    not_broadcastable = (
        ([2, 3], [2, 3, 4], [0, 2]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # 5. TypeError: broadcast_in_dim shape must have every element be nonnegative, got (-1, 2, 3).
    negative_shape = (
        ([2, 3], [2, 3, -1], [0, 1]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # TODO add exceptions for not_broadcastable, negative output shape
    error_cases = [
        missing_axis_in_bcast_dims,
        fewer_dims_in_output_shape,
        out_of_bounds_broadcast_dimensions,
        # not_broadcastable,
        # negative_shape,
    ]
    for es in error_cases:
        ex_case, ex_type, ex_str = es
        input_shape, output_shape, bcast_dims = ex_case
        input_tensor = make_arg(input_shape)
        if op.name == "broadcast_in_dim_symbolic":
            bcast_shaped_tensor = make_arg(output_shape)
            yield SampleInput(
                input_tensor, bcast_shaped_tensor, bcast_dims
            ), ex_type, ex_str
        else:
            yield SampleInput(input_tensor, output_shape, bcast_dims), ex_type, ex_str

def cat_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # concatenating tensors along singleton, broadcast dimensions is unsupported by nvFuser.
    # https://github.com/NVIDIA/Fuser/issues/224
    # shapes, dim
    cases = [
        ([(3,)], 0),  # single tensor provided
        # 1D
        ([(2,), (3,)], 0),
        ([(2,), (4,)], 0),
        ([(0,), (2,)], 0),
        ([(0,), (2,)], -1),
        ([(2, 3), (2, 4)], 1),
        ([(2, 3), (2, 4), (2, 5)], 1),
    ]

    for shapes, dim in cases:
        yield SampleInput([make_arg(s) for s in shapes], dim)

def cat_error_generator(op, dtype=torch.float32, requires_grad: bool = False, **kwargs):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )
    # shapes, dim, exception type, exception string
    empty_input_tensors = (
        ([], 0),
        RuntimeError,
        "Attempting to concatenate empty list of tensors",
    )
    positive_dim = (([(1,), (2,)], 1), RuntimeError, "Invalid dimension to cat")
    negative_dim = (([(2,), (2,)], -2), RuntimeError, "Invalid dimension to cat")
    # All tensors must have same number of dimension"
    ndims_mismatch = (
        ([(2,), (2, 3)], 0),
        RuntimeError,
        "Unexpected number of dimensions",
    )
    # All tensors must have same shape except for the cat dimension
    shape_mismatch = (
        ([(2, 3), (4, 5)], 0),
        RuntimeError,
        "a conflict was found with 2 different sizes",
    )

    error_cases = [
        empty_input_tensors,
        positive_dim,
        negative_dim,
        ndims_mismatch,
        shape_mismatch,
    ]

    for case, ex_type, ex_str in error_cases:
        shapes, dim = case
        yield SampleInput([make_arg(s) for s in shapes], dim), ex_type, ex_str

def define_tensor_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    "define_tensor",
    [](FusionDefinition& self,
        std::vector<int64_t>& sizes,
        std::vector<int64_t>& strides,
        PrimDataType dtype = DataType::Float,
        bool static_sizes = false,
        bool is_cpu = false) -> Tensor {
    ---
    "define_tensor",
    [](FusionDefinition& self,
        std::vector<int64_t>& shape,
        std::vector<std::optional<bool>>& contiguity,
        PrimDataType dtype = DataType::Float,
        bool is_cpu = false) -> Tensor {
    """

    check_size_contiguity_match = ErrorSample(
        {
            "shape": [-1, -1],
            "contiguity": [True, True, True],
            "dtype": DataType.Float,
        },
        "Length of contiguity argument (3) must match that of shape argument (2)",
    )

    check_empty_tensor_size = ErrorSample(
        {"shape": [], "contiguity": []},
        "Empty tensor is unsupported.",
    )

    check_max_tensor_size = ErrorSample(
        {
            "shape": [-1 for _ in range(MAX_TENSOR_DIMS + 1)],
            "contiguity": [True for _ in range(MAX_TENSOR_DIMS + 1)],
        },
        "The specified tensor dimensionality exceeds the max tensor size for nvfuser.",
    )

    check_above_size_range = ErrorSample(
        {"shape": [INT64_MAX + 1], "contiguity": [True]},
        "define_tensor(): incompatible function arguments",
        TypeError,
    )

    check_below_size_range = ErrorSample(
        {"shape": [MINIMUM_SYMBOLIC_SIZE - 1], "contiguity": [True]},
        "The value -2 at index 0 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)",
    )

    check_contiguity_unknown_values = ErrorSample(
        {"shape": [10], "contiguity": [-1]},
        "define_tensor(): incompatible function arguments.",
        TypeError,
    )

    check_shape_unknown_dtypes = ErrorSample(
        {"shape": [10.0], "contiguity": [True]},
        "define_tensor(): incompatible function arguments.",
        TypeError,
    )

    check_stride_order_duplicate = ErrorSample(
        {
            "shape": [-1, -1, -1],
            "contiguity": [True, True, True],
            "stride_order": [0, 1, 1],
        },
        "duplicated stride_order entries",
        RuntimeError,
    )

    check_stride_order_out_of_range = ErrorSample(
        {
            "shape": [-1, -1, -1],
            "contiguity": [True, True, True],
            "stride_order": [0, 1, 5],
        },
        "stride_order argument is out of range",
        RuntimeError,
    )

    check_stride_order_out_of_negative_range = ErrorSample(
        {
            "shape": [-1, -1, -1],
            "contiguity": [True, True, True],
            "stride_order": [0, 1, -4],
        },
        "stride_order argument is out of range",
        RuntimeError,
    )

    # TODO: Fix empty and maximum tensor dimensionality error checks.
    # TODO: Add invalid argument checks for contiguity.
    error_cases = [
        check_size_contiguity_match,
        # check_empty_tensor_size,
        # check_max_tensor_size,
        check_above_size_range,
        check_below_size_range,
        # check_contiguity_unknown_values,
        check_shape_unknown_dtypes,
    ]

    input_tensor = make_tensor(
        (10, 10), device="cuda", dtype=dtype, requires_grad=requires_grad
    )
    for es in error_cases:
        yield SampleInput(input_tensor, **es.kwargs), es.ex_type, es.ex_str

def define_vector_constant_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    "define_vector",
    [](FusionDefinition& self, py::list& values) -> Vector {
    """

    check_above_size_range = ErrorSample(
        {"values": [INT64_MAX + 1]},
        "define_vector(): incompatible function arguments",
        TypeError,
    )

    check_below_size_range = ErrorSample(
        {"values": [MINIMUM_SYMBOLIC_SIZE - 1]},
        "The value -2 at index 0 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)",
    )

    error_cases = [
        # FIXME: The above_size_range case gives a non-sensical error message.
        # "Unable to cast Python instance to C++ type (#define PYBIND11_DETAILED_ER"
        # check_above_size_range,
        check_below_size_range,
    ]

    for es in error_cases:
        yield SampleInput(**es.kwargs), es.ex_type, es.ex_str

def _special_value_binary_generator(
    lhs_generator_fn, rhs_generator_fn, dtype, requires_grad
):
    lhs_vals, rhs_vals = zip(*itertools.product(lhs_generator_fn, rhs_generator_fn))
    lhs = torch.tensor(
        lhs_vals, device="cuda", dtype=dtype, requires_grad=requires_grad
    )
    rhs = torch.tensor(
        rhs_vals, device="cuda", dtype=dtype, requires_grad=requires_grad
    )
    return SampleInput(lhs, rhs)

def elementwise_binary_generator(
    op: OpInfo,
    dtype: torch.dtype,
    requires_grad: bool = False,
    *,
    supports_numbers: bool = True,
    enable_broadcast_testing: bool = True,
    enable_extremal_value_testing: bool = True,
    enable_large_value_testing: bool = True,
    enable_small_value_testing: bool = True,
    **kwargs,
):
    low = None if op.domain.low is None else max(-9, op.domain.low)
    high = None if op.domain.high is None else min(9, op.domain.high)
    make_arg = partial(
        make_tensor,
        device="cuda",
        dtype=dtype,
        low=low,
        high=high,
        requires_grad=requires_grad,
        **kwargs,
    )

    shapes = (
        (0, 2, 1),
        (5, 0, 3),
        (),
        (11,),
        (4, 4),
        (1024, 1024),
        (64, 64, 64),
    )

    # Typical inputs
    for shape in shapes:
        yield SampleInput(make_arg(shape), make_arg(shape))
        yield SampleInput(
            make_arg(shape, noncontiguous=True), make_arg(shape, noncontiguous=True)
        )

    if enable_broadcast_testing:
        broadcast_shapes = (
            ((1,), ()),
            ((2,), ()),
            ((1,), (2,)),
            ((2, 1), (2,)),
            ((1, 2), (2,)),
            ((3, 2), (2,)),
            ((1, 3, 2), (2,)),
            ((1, 3, 2), (3, 2)),
            ((3, 1, 2), (3, 2)),
            ((2, 3, 2), ()),
            ((3, 1, 2), (1, 3, 2)),
        )
        for lhs_shape, rhs_shape in broadcast_shapes:
            yield SampleInput(make_arg(lhs_shape), make_arg(rhs_shape))
            yield SampleInput(
                make_arg(lhs_shape, noncontiguous=True),
                make_arg(rhs_shape, noncontiguous=True),
            )

    # Create filtered special inputs for this operation's domain
    def _filter_lhs_domain(values):
        return [v for v in values if is_within_domain(op.domain, v)]

    def _filter_rhs_domain(values):
        # NOTE: Check exclude_zero flag to avoid undefined behavior such as ZeroDivisionError: division by zero
        exclude_zero = kwargs.get("exclude_zero", False)
        return [v for v in values if is_within_domain(op.domain, v, exclude_zero)]

    if (
        enable_large_value_testing
        and dtype != torch.bool
        and dtype not in complex_dtypes
    ):
        lhs_large_values = _filter_lhs_domain(_large_values(dtype))
        rhs_large_values = _filter_rhs_domain(_large_values(dtype))
        yield _special_value_binary_generator(
            lhs_large_values, rhs_large_values, dtype, requires_grad
        )

    if enable_small_value_testing and dtype != torch.bool:
        lhs_small_values = _filter_lhs_domain(_small_values(dtype))
        rhs_small_values = _filter_rhs_domain(_small_values(dtype))
        yield _special_value_binary_generator(
            lhs_small_values, rhs_small_values, dtype, requires_grad
        )

    if enable_extremal_value_testing and dtype in float_complex_dtypes:
        lhs_extremal_values = _filter_lhs_domain(_extremal_values(dtype))
        rhs_extremal_values = _filter_rhs_domain(_extremal_values(dtype))
        yield _special_value_binary_generator(
            lhs_extremal_values, rhs_extremal_values, dtype, requires_grad
        )

        # Test interactions between extreme and normal values
        make_cuda_tensor = partial(
            torch.tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
        )
        rhs_normal = [random.uniform(-10, 10) for _ in range(len(lhs_extremal_values))]
        lhs_normal = [random.uniform(-10, 10) for _ in range(len(rhs_extremal_values))]
        yield SampleInput(
            make_cuda_tensor(lhs_extremal_values), make_cuda_tensor(rhs_normal)
        )
        yield SampleInput(
            make_cuda_tensor(lhs_normal), make_cuda_tensor(rhs_extremal_values)
        )

def _elementwise_binary_torch(op):
    @wraps(op)
    def _fn(x, y):
        if isinstance(x, torch.Tensor) or isinstance(y, torch.Tensor):
            return op(x, y)
        return op(torch.tensor(x), torch.tensor(y)).item()

    return _fn

def elementwise_unary_generator(
    op: OpInfo,
    dtype: torch.dtype,
    requires_grad: bool = False,
    *,
    supports_numbers: bool = True,
    enable_extremal_value_testing: bool = True,
    enable_large_value_testing: bool = True,
    enable_small_value_testing: bool = True,
    **kwargs,
):
    low = None if op.domain.low is None else max(-9, op.domain.low)
    high = None if op.domain.high is None else min(9, op.domain.high)
    make_arg = partial(
        make_tensor,
        device="cuda",
        dtype=dtype,
# ... truncated ...


    is_floating_dtype,
    float_complex_dtypes,
    complex_dtypes,
)
import sys

if "nvfuser" in sys.modules:
    from nvfuser import DataType
    from nvfuser.pytorch_utils import torch_dtype_to_nvfuser_dtype
else:
    from nvfuser_direct import DataType
    from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

MINIMUM_SYMBOLIC_SIZE = -1
INT64_MAX = 2**63 - 1
MAX_TENSOR_DIMS = 8

# Determine if a number is with desired Domain [low, high)
# The domain is half-open. The lower limit is inclusive while the upper limit is exclusive.
def is_within_domain(domain: Domain, a: Number, exclude_zero: bool = False):
    # comparison operators are not defined for complex numbers
    if isinstance(a, complex):
        return True

    if domain.low is not None and domain.low > a:
        return False

    if domain.high is not None and a >= domain.high:
        return False

    if exclude_zero and a == 0:
        return False

    return True

def _extremal_values(dtype: torch.dtype):
    _float_vals = (float("inf"), float("-inf"), float("nan"))
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )
    _int16_vals = (-32768, 32767)
    _int32_vals = (-2147483648, 2147483647)
    _int64_vals = (-9223372036854775808, 9223372036854775807)

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype is torch.int16:
        return _int16_vals
    elif dtype is torch.int32:
        return _int32_vals
    elif dtype is torch.int64:


# Determine if a number is with desired Domain [low, high)
# The domain is half-open. The lower limit is inclusive while the upper limit is exclusive.
def is_within_domain(domain: Domain, a: Number, exclude_zero: bool = False):
    # comparison operators are not defined for complex numbers
    if isinstance(a, complex):
        return True

    if domain.low is not None and domain.low > a:
        return False

    if domain.high is not None and a >= domain.high:
        return False

    if exclude_zero and a == 0:
        return False

    return True

def _extremal_values(dtype: torch.dtype):
    _float_vals = (float("inf"), float("-inf"), float("nan"))
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )
    _int16_vals = (-32768, 32767)
    _int32_vals = (-2147483648, 2147483647)
    _int64_vals = (-9223372036854775808, 9223372036854775807)

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype is torch.int16:
        return _int16_vals
    elif dtype is torch.int32:
        return _int32_vals
    elif dtype is torch.int64:
        return _int64_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _large_values(dtype: torch.dtype):
    _int_vals = (-1113, 1113, -10701, 10701)
    _float16_vals = (-501, 501, -1001.2, 1001.2, -13437.7, 13437.7)
    _float_vals = _float16_vals + (-4988429.2, 4988429.2, -1e20, 1e20)
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype == torch.float16:
        return _float16_vals
    elif dtype in (torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        print(_complex_vals)
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:


        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )
    _int16_vals = (-32768, 32767)
    _int32_vals = (-2147483648, 2147483647)
    _int64_vals = (-9223372036854775808, 9223372036854775807)

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype is torch.int16:
        return _int16_vals
    elif dtype is torch.int32:
        return _int32_vals
    elif dtype is torch.int64:
        return _int64_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _large_values(dtype: torch.dtype):
    _int_vals = (-1113, 1113, -10701, 10701)
    _float16_vals = (-501, 501, -1001.2, 1001.2, -13437.7, 13437.7)
    _float_vals = _float16_vals + (-4988429.2, 4988429.2, -1e20, 1e20)
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype == torch.float16:
        return _float16_vals
    elif dtype in (torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        print(_complex_vals)
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _small_values(dtype: torch.dtype):
    eps = 1e-5
    _int_vals = (0, -1, 1, -55, 55, -127, 127, -128)
    _float_vals = (
        0.0,
        -0.0,
        -1e-3,
        1e-3,
        -0.25,
        0.25,
        -1.0,
        1.0,
        -math.e / 2.0,
        math.e / 2.0,
        -math.e + eps,
        math.e - eps,
        -math.e,
        math.e,


    _int_vals = (-1113, 1113, -10701, 10701)
    _float16_vals = (-501, 501, -1001.2, 1001.2, -13437.7, 13437.7)
    _float_vals = _float16_vals + (-4988429.2, 4988429.2, -1e20, 1e20)
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype == torch.float16:
        return _float16_vals
    elif dtype in (torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        print(_complex_vals)
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def _small_values(dtype: torch.dtype):
    eps = 1e-5
    _int_vals = (0, -1, 1, -55, 55, -127, 127, -128)
    _float_vals = (
        0.0,
        -0.0,
        -1e-3,
        1e-3,
        -0.25,
        0.25,
        -1.0,
        1.0,
        -math.e / 2.0,
        math.e / 2.0,
        -math.e + eps,
        math.e - eps,
        -math.e,
        math.e,
        -math.e - eps,
        math.e + eps,
    )
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def broadcast_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast(operand, sizes)
    # add new dimensions to left-hand-side of tensor
    # dims = tuple(range(len(sizes), len(sizes) + np.ndim(operand)))
    # return broadcast_in_dim(operand, tuple(sizes) + np.shape(operand), dims)

    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    fewer_original_axes = (
        ([2, 3], [True, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )


        math.e - eps,
        -math.e,
        math.e,
        -math.e - eps,
        math.e + eps,
    )
    _complex_vals = tuple(
        complex(*x) for x in itertools.product(_float_vals, _float_vals)
    )

    if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        return _float_vals
    elif dtype in (torch.complex64, torch.complex128):
        return _complex_vals
    elif dtype in (torch.int16, torch.int32, torch.int64):
        return _int_vals
    else:
        raise ValueError(f"Unsupported dtype --- {dtype}")

def broadcast_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast(operand, sizes)
    # add new dimensions to left-hand-side of tensor
    # dims = tuple(range(len(sizes), len(sizes) + np.ndim(operand)))
    # return broadcast_in_dim(operand, tuple(sizes) + np.shape(operand), dims)

    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    fewer_original_axes = (
        ([2, 3], [True, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    greater_original_axes = (
        ([2, 3], [True, False, False, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    error_cases = [
        fewer_original_axes,
        greater_original_axes,
    ]
    for es in error_cases:
        ex_case, ex_type, ex_str = es
        input_shape, bcast_dims = ex_case
        input_tensor = make_arg(input_shape)
        yield SampleInput(input_tensor, bcast_dims), ex_type, ex_str

def broadcast_in_dim_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # The first 5 test cases below are taken from JAX's broadcast_in_dim tests
    #   https://github.com/google/jax/blob/main/tests/lax_test.py#L1171
    # input shape, output shape, bcast_dims
    cases = (
        ([2], [2, 2], [0]),
        ([2], [2, 2], [1]),
        ([2], [2, 3], [0]),
        ([], [2, 3], []),
        ([1], [2, 3], [1]),
        ((4, 6, 3, 1), (5, 4, 7, 6, 3, 6, 6), (1, 3, 4, 5)),
    )


        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    greater_original_axes = (
        ([2, 3], [True, False, False, False]),
        RuntimeError,
        "Invalid broadcast, number of false entries in is_broadcast_dim expected to be",
    )

    error_cases = [
        fewer_original_axes,
        greater_original_axes,
    ]
    for es in error_cases:
        ex_case, ex_type, ex_str = es
        input_shape, bcast_dims = ex_case
        input_tensor = make_arg(input_shape)
        yield SampleInput(input_tensor, bcast_dims), ex_type, ex_str

def broadcast_in_dim_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # The first 5 test cases below are taken from JAX's broadcast_in_dim tests
    #   https://github.com/google/jax/blob/main/tests/lax_test.py#L1171
    # input shape, output shape, bcast_dims
    cases = (
        ([2], [2, 2], [0]),
        ([2], [2, 2], [1]),
        ([2], [2, 3], [0]),
        ([], [2, 3], []),
        ([1], [2, 3], [1]),
        ((4, 6, 3, 1), (5, 4, 7, 6, 3, 6, 6), (1, 3, 4, 5)),
    )

    for input_shape, output_shape, bcast_dims in cases:
        input_tensor = make_arg(input_shape)
        if op.name == "broadcast_in_dim_symbolic":
            bcast_shaped_tensor = make_arg(output_shape)
            yield SampleInput(input_tensor, bcast_shaped_tensor, bcast_dims)
        else:
            yield SampleInput(input_tensor, output_shape, bcast_dims)

def broadcast_in_dim_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # 1. Every dimension in the input tensor must be used in broadcast_dimensions.
    missing_axis_in_bcast_dims = (
        ([2, 2], [2, 2, 3], [0]),
        RuntimeError,
        "The broadcast dimensions should match the input dimensions.",
    )

    # 2. New shape has weakly more dimentions than the original tensor.
    fewer_dims_in_output_shape = (
        ([2, 2], [2], [0]),


    #   https://github.com/google/jax/blob/main/tests/lax_test.py#L1171
    # input shape, output shape, bcast_dims
    cases = (
        ([2], [2, 2], [0]),
        ([2], [2, 2], [1]),
        ([2], [2, 3], [0]),
        ([], [2, 3], []),
        ([1], [2, 3], [1]),
        ((4, 6, 3, 1), (5, 4, 7, 6, 3, 6, 6), (1, 3, 4, 5)),
    )

    for input_shape, output_shape, bcast_dims in cases:
        input_tensor = make_arg(input_shape)
        if op.name == "broadcast_in_dim_symbolic":
            bcast_shaped_tensor = make_arg(output_shape)
            yield SampleInput(input_tensor, bcast_shaped_tensor, bcast_dims)
        else:
            yield SampleInput(input_tensor, output_shape, bcast_dims)

def broadcast_in_dim_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # 1. Every dimension in the input tensor must be used in broadcast_dimensions.
    missing_axis_in_bcast_dims = (
        ([2, 2], [2, 2, 3], [0]),
        RuntimeError,
        "The broadcast dimensions should match the input dimensions.",
    )

    # 2. New shape has weakly more dimentions than the original tensor.
    fewer_dims_in_output_shape = (
        ([2, 2], [2], [0]),
        RuntimeError,
        "The new shape is expected to be greater-then-or-equal to the input",
    )

    # 3. Each broadcast dimension is within the new shape.
    out_of_bounds_broadcast_dimensions = (
        ([2, 2], [2, 2], [0, 2]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # 4. The original tensor is not broadcastable to desired shape.
    # tensor.shape[idx] == 1 or tensor.shape[idx] == output_shape[new_idx]
    #
    # Jax Exception:
    # TypeError: broadcast_in_dim operand dimension sizes must either be 1,
    # or be equal to their corresponding dimensions in the target broadcast shape;
    # got operand of shape (2, 3), target broadcast shape (2, 3, 4), broadcast_dimensions (0, 2)
    not_broadcastable = (
        ([2, 3], [2, 3, 4], [0, 2]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # 5. TypeError: broadcast_in_dim shape must have every element be nonnegative, got (-1, 2, 3).
    negative_shape = (
        ([2, 3], [2, 3, -1], [0, 1]),
        RuntimeError,
        "Invalid broadcast_dims value.",
    )

    # TODO add exceptions for not_broadcastable, negative output shape
    error_cases = [
        missing_axis_in_bcast_dims,
        fewer_dims_in_output_shape,
        out_of_bounds_broadcast_dimensions,
        # not_broadcastable,
        # negative_shape,
    ]
    for es in error_cases:
        ex_case, ex_type, ex_str = es
        input_shape, output_shape, bcast_dims = ex_case
        input_tensor = make_arg(input_shape)
        if op.name == "broadcast_in_dim_symbolic":
            bcast_shaped_tensor = make_arg(output_shape)
            yield SampleInput(
                input_tensor, bcast_shaped_tensor, bcast_dims
            ), ex_type, ex_str
        else:
            yield SampleInput(input_tensor, output_shape, bcast_dims), ex_type, ex_str

def cat_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    make_arg = partial(
        make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad
    )

    # concatenating tensors along singleton, broadcast dimensions is unsupported by nvFuser.
    # https://github.com/NVIDIA/Fuser/issues/224
    # shapes, dim
    cases = [
        ([(3,)], 0),  # single tensor provided
        # 1D
        ([(2,), (3,)], 0),
        ([(2,), (4,)], 0),
        ([(0,), (2,)], 0),
        ([(0,), (2,)], -1),
        ([(2, 3), (2, 4)], 1),