        for offset in offsets:
            yield SampleInput(*element.args, offset)

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

    invalid_shapes = (
        (),
        (4,),
    )

    for shape in invalid_shapes:
        yield SampleInput(
            make_arg(shape),
        ), RuntimeError, f"input tensor for triu must have 2 or more dims, but got {len(shape)} dims"

def cumsum_input_generator(op: OpInfo, dtype: torch.dtype, requires_grad: bool = False):
    make_arg = partial(
        make_tensor,
        device="cuda",
        dtype=dtype,
        requires_grad=requires_grad,
        low=-2,
        high=3,
    )

    # shape, dim
    cases = (
        ((4, 4), 0),
        ((5,), 0),
        ((8, 1, 6), 1),
        ((8, 7, 5, 1), 2),
        ((8, 7, 5, 1), -1),
        ((8, 7, 5, 1), -2),
    )

    for shape, dim in cases:
        yield SampleInput(make_arg(shape), dim)

def cumsum_error_generator(
    op: OpInfo,
    dtype: torch.dtype,
    requires_grad: bool = False,
):
    """
    Generate test cases that should produce errors for cumsum operation.

    Yields:
        Tuples of (SampleInput, expected_exception_type, error_message_pattern)
    """
    make_arg = partial(
        make_tensor,
        device="cuda",
        dtype=dtype,
        requires_grad=requires_grad,
        low=-2,
        high=3,


        dtype=dtype,
        requires_grad=requires_grad,
        low=-2,
        high=3,
    )

    # shape, dim
    cases = (
        ((4, 4), 0),
        ((5,), 0),
        ((8, 1, 6), 1),
        ((8, 7, 5, 1), 2),
        ((8, 7, 5, 1), -1),
        ((8, 7, 5, 1), -2),
    )

    for shape, dim in cases:
        yield SampleInput(make_arg(shape), dim)

def cumsum_error_generator(
    op: OpInfo,
    dtype: torch.dtype,
    requires_grad: bool = False,
):
    """
    Generate test cases that should produce errors for cumsum operation.

    Yields:
        Tuples of (SampleInput, expected_exception_type, error_message_pattern)
    """
    make_arg = partial(
        make_tensor,
        device="cuda",
        dtype=dtype,
        requires_grad=requires_grad,
        low=-2,
        high=3,
    )

    # Invalid dimension: out of bounds positive
    input_shape = (4, 4)
    yield SampleInput(
        make_arg(input_shape), 2
    ), RuntimeError, "Tried to access out of boundary index"

    # Invalid dimension: out of bounds negative
    yield SampleInput(
        make_arg(input_shape), -3
    ), RuntimeError, "Tried to access out of boundary index"

    # Invalid dimension type: float instead of int
    yield SampleInput(
        make_arg(input_shape), 0.5
    ), TypeError, "cumsum(): incompatible function arguments"

    # Invalid input: 0-dim tensor (scalar)
    yield SampleInput(
        make_arg(()), 0
    ), RuntimeError, "Tried to access out of boundary index 0. total index: 0"

def grouped_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for grouped matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """
    make_arg = partial(
        make_tensor,
        dtype=dtype,
        device="cuda",
        low=None,
        high=None,
        requires_grad=requires_grad,


    yield SampleInput(
        make_arg(input_shape), 2
    ), RuntimeError, "Tried to access out of boundary index"

    # Invalid dimension: out of bounds negative
    yield SampleInput(
        make_arg(input_shape), -3
    ), RuntimeError, "Tried to access out of boundary index"

    # Invalid dimension type: float instead of int
    yield SampleInput(
        make_arg(input_shape), 0.5
    ), TypeError, "cumsum(): incompatible function arguments"

    # Invalid input: 0-dim tensor (scalar)
    yield SampleInput(
        make_arg(()), 0
    ), RuntimeError, "Tried to access out of boundary index 0. total index: 0"

def grouped_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for grouped matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """
    make_arg = partial(
        make_tensor,
        dtype=dtype,
        device="cuda",
        low=None,
        high=None,
        requires_grad=requires_grad,
    )

    def make_index(extent, num_groups):
        group_size = extent // num_groups
        return torch.arange(
            group_size,
            group_size * g + 1,
            group_size,
            device="cuda",
            dtype=torch.int32,
            requires_grad=False,
        )

    # TODO: expand the test when kernel restrictions are lifted
    # Test various group sizes and matrix dimensions
    configs = (
        (4, 128, 256, 64),
        (2, 32, 32, 32),
    )

    for config in configs:
        g, m, k, n = config

        # case 1: 2d x 2d
        mat1 = make_arg((m, k))
        mat2 = make_arg((k, n))
        offsets = make_index(k, g)
        yield SampleInput(mat1, mat2, offsets)
        # case 3: 2d x 3d
        mat1 = make_arg((m, k))
        mat2 = make_arg((g, k, n))
        offsets = make_index(m, g)
        yield SampleInput(mat1, mat2, offsets)
        # case 1: 3d x 2d
        mat1 = make_arg((g, m, k))
        mat2 = make_arg((k, n))
        offsets = make_index(n, g)
        yield SampleInput(mat1, mat2, offsets)

def scaled_grouped_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for scaled grouped matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """

    # TODO: enable mxfp8 test when backend supports it.
    make_arg = partial(
        make_tensor,
        dtype=dtype,
        device="cuda",
        low=None,


    for config in configs:
        g, m, k, n = config

        # case 1: 2d x 2d
        mat1 = make_arg((m, k))
        mat2 = make_arg((k, n))
        offsets = make_index(k, g)
        yield SampleInput(mat1, mat2, offsets)
        # case 3: 2d x 3d
        mat1 = make_arg((m, k))
        mat2 = make_arg((g, k, n))
        offsets = make_index(m, g)
        yield SampleInput(mat1, mat2, offsets)
        # case 1: 3d x 2d
        mat1 = make_arg((g, m, k))
        mat2 = make_arg((k, n))
        offsets = make_index(n, g)
        yield SampleInput(mat1, mat2, offsets)

def scaled_grouped_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for scaled grouped matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """

    # TODO: enable mxfp8 test when backend supports it.
    make_arg = partial(
        make_tensor,
        dtype=dtype,
        device="cuda",
        low=None,
        high=None,
        requires_grad=requires_grad,
    )

    make_scale_factor = partial(
        make_tensor,
        dtype=torch.float32,
        device="cuda",
        low=None,
        high=None,
        requires_grad=False,
    )

    def make_index(extent, num_groups):
        group_size = extent // num_groups
        return torch.arange(
            group_size,
            group_size * g + 1,
            group_size,
            device="cuda",
            dtype=torch.int32,
            requires_grad=False,
        )

    # TODO: expand the test when fallback kernel restrictions are lifted
    #       currently only bf16 output is supported.
    #       there are also restrictions on the input/output shapes.
    # Test various group sizes and matrix dimensions
    # configs: list(g, m, k, n, output_dtype)
    configs = (
        (4, 128, 256, 64, torch.bfloat16),
        (2, 32, 32, 32, torch.bfloat16),
    )

    # TODO: Enable mxfp8 test when backend supports it.
    for config in configs:
        g, m, k, n, dtype = config
        # case 1: 2d x 2d
        mat1 = make_arg((m, k))
        mat2 = make_arg((k, n))
        scale1 = make_scale_factor((g, m, 1))
        scale2 = make_scale_factor((g, 1, n))
        offsets = make_index(k, g)
        yield SampleInput(mat1, mat2, offsets, scale1, scale2, None, None, None, dtype)
        # case 3: 2d x 3d
        mat1 = make_arg((m, k))
        mat2 = make_arg((g, k, n))
        scale1 = make_scale_factor((m, 1))
        scale2 = make_scale_factor((g, 1, n))
        offsets = make_index(m, g)
        yield SampleInput(mat1, mat2, offsets, scale1, scale2, None, None, None, dtype)
        # case 1: 3d x 2d
        mat1 = make_arg((g, m, k))
        mat2 = make_arg((k, n))
        offsets = make_index(n, g)
        scale1 = make_scale_factor((g, m, 1))
        scale2 = make_scale_factor((1, n))
        yield SampleInput(mat1, mat2, offsets, scale1, scale2, None, None, None, dtype)

###############################################################################
#
# NOTE: quantization taken from torch testing. Since this is using aten backend
#
###############################################################################
# largest power of 2 representable in `torch.float8_e4m3fn`
F8E4M3_LARGEST_POW2 = 8
# max value of `torch.float8_e4m3fn` (448)
F8E4M3_MAX_VAL = torch.finfo(torch.float8_e4m3fn).max
# exponent bias of `torch.float8_e8m0fnu`
F8E8M0_EXP_BIAS = 127

def data_to_mxfp8_scale(x, block_size):
    # simple implementation of https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
    # section 6.3, not all edge cases (such as NaN) are handled/tested
    orig_shape = x.shape
    x = x.reshape(-1, block_size)


        mat2 = make_arg((k, n))
        offsets = make_index(n, g)
        scale1 = make_scale_factor((g, m, 1))
        scale2 = make_scale_factor((1, n))
        yield SampleInput(mat1, mat2, offsets, scale1, scale2, None, None, None, dtype)

###############################################################################
#
# NOTE: quantization taken from torch testing. Since this is using aten backend
#
###############################################################################
# largest power of 2 representable in `torch.float8_e4m3fn`
F8E4M3_LARGEST_POW2 = 8
# max value of `torch.float8_e4m3fn` (448)
F8E4M3_MAX_VAL = torch.finfo(torch.float8_e4m3fn).max
# exponent bias of `torch.float8_e8m0fnu`
F8E8M0_EXP_BIAS = 127

def data_to_mxfp8_scale(x, block_size):
    # simple implementation of https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
    # section 6.3, not all edge cases (such as NaN) are handled/tested
    orig_shape = x.shape
    x = x.reshape(-1, block_size)
    max_abs = torch.amax(torch.abs(x), 1)
    largest_p2_lt_max_abs = torch.floor(torch.log2(max_abs))
    scale_e8m0_unbiased = largest_p2_lt_max_abs - F8E4M3_LARGEST_POW2
    scale_e8m0_unbiased = torch.clamp(
        scale_e8m0_unbiased, -1 * F8E8M0_EXP_BIAS, F8E8M0_EXP_BIAS
    )
    scale_e8m0_biased = scale_e8m0_unbiased + F8E8M0_EXP_BIAS
    scale_e8m0_biased = scale_e8m0_biased.to(torch.uint8)
    scale_e8m0_biased = scale_e8m0_biased.view(torch.float8_e8m0fnu)
    return scale_e8m0_biased.reshape(orig_shape[0], -1)

def data_to_mxfp8(x, block_size):
    x_scale = data_to_mxfp8_scale(x, block_size)
    K = x.shape[1]
    max_val = F8E4M3_MAX_VAL
    min_val = -1 * max_val
    x = (x.reshape(-1, block_size) / x_scale.reshape(-1, 1).float()).reshape(-1, K)
    x = x.clamp(min=min_val, max=max_val).to(torch.float8_e4m3fn)
    return x, x_scale

def scaled_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for scaled matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation


F8E8M0_EXP_BIAS = 127

def data_to_mxfp8_scale(x, block_size):
    # simple implementation of https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
    # section 6.3, not all edge cases (such as NaN) are handled/tested
    orig_shape = x.shape
    x = x.reshape(-1, block_size)
    max_abs = torch.amax(torch.abs(x), 1)
    largest_p2_lt_max_abs = torch.floor(torch.log2(max_abs))
    scale_e8m0_unbiased = largest_p2_lt_max_abs - F8E4M3_LARGEST_POW2
    scale_e8m0_unbiased = torch.clamp(
        scale_e8m0_unbiased, -1 * F8E8M0_EXP_BIAS, F8E8M0_EXP_BIAS
    )
    scale_e8m0_biased = scale_e8m0_unbiased + F8E8M0_EXP_BIAS
    scale_e8m0_biased = scale_e8m0_biased.to(torch.uint8)
    scale_e8m0_biased = scale_e8m0_biased.view(torch.float8_e8m0fnu)
    return scale_e8m0_biased.reshape(orig_shape[0], -1)

def data_to_mxfp8(x, block_size):
    x_scale = data_to_mxfp8_scale(x, block_size)
    K = x.shape[1]
    max_val = F8E4M3_MAX_VAL
    min_val = -1 * max_val
    x = (x.reshape(-1, block_size) / x_scale.reshape(-1, 1).float()).reshape(-1, K)
    x = x.clamp(min=min_val, max=max_val).to(torch.float8_e4m3fn)
    return x, x_scale

def scaled_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for scaled matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """

    # TODO: enable mxfp8 test when backend supports it.
    make_arg = partial(
        make_tensor,
        dtype=torch.float32,
        device="cuda",
        low=None,


    scale_e8m0_unbiased = largest_p2_lt_max_abs - F8E4M3_LARGEST_POW2
    scale_e8m0_unbiased = torch.clamp(
        scale_e8m0_unbiased, -1 * F8E8M0_EXP_BIAS, F8E8M0_EXP_BIAS
    )
    scale_e8m0_biased = scale_e8m0_unbiased + F8E8M0_EXP_BIAS
    scale_e8m0_biased = scale_e8m0_biased.to(torch.uint8)
    scale_e8m0_biased = scale_e8m0_biased.view(torch.float8_e8m0fnu)
    return scale_e8m0_biased.reshape(orig_shape[0], -1)

def data_to_mxfp8(x, block_size):
    x_scale = data_to_mxfp8_scale(x, block_size)
    K = x.shape[1]
    max_val = F8E4M3_MAX_VAL
    min_val = -1 * max_val
    x = (x.reshape(-1, block_size) / x_scale.reshape(-1, 1).float()).reshape(-1, K)
    x = x.clamp(min=min_val, max=max_val).to(torch.float8_e4m3fn)
    return x, x_scale

def scaled_mm_input_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    """
    Generate valid test cases for scaled matrix multiplication.

    Args:
        op: OpInfo object for the bmm operation
        dtype: Data type for test tensors
        requires_grad: Whether tensors should require gradients
    """

    # TODO: enable mxfp8 test when backend supports it.
    make_arg = partial(
        make_tensor,
        dtype=torch.float32,
        device="cuda",
        low=None,
        high=None,
        requires_grad=requires_grad,
    )

    # TODO: expand the test when fallback kernel restrictions are lifted
    #       currently only bf16 output is supported.
    #       there are also restrictions on the input/output shapes.
    # Test various group sizes and matrix dimensions
    # configs: list(m, k, n, output_dtype)
    configs = (
        (128, 256, 512, torch.bfloat16),
        (128, 128, 128, torch.bfloat16),
    )

    # TODO: support nvfp4
    assert dtype == torch.float8_e4m3fn
    quantization = partial(data_to_mxfp8, block_size=32)

    for config in configs:
        m, k, n, dtype = config
        mat1_ref = make_arg((m, k))
        mat2_ref = make_arg((n, k))
        mat1, scale1 = quantization(mat1_ref)
        mat2, scale2 = quantization(mat2_ref)
        yield SampleInput(mat1, mat2.t(), scale1, scale2, None, None, None, dtype)


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


            ((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:


                device="cuda",
                dtype=dtype,
                requires_grad=requires_grad,
            )
        )

    if enable_extremal_value_testing and dtype in float_complex_dtypes:
        filtered_extremal_values = _filter_domain(_extremal_values(dtype))
        yield SampleInput(
            torch.tensor(
                filtered_extremal_values,
                device="cuda",
                dtype=dtype,
                requires_grad=requires_grad,
            )
        )

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

    return _fn

def full_error_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # torch.full(size, fill_value, dtype=None)
    # Error: Trying to create tensor with negative dimension
    negative_input_shape = [2, -2]
    yield SampleInput(
        negative_input_shape, make_number(dtype), dtype
    ), RuntimeError, "The value -2 at index 1 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)."

def scatter_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # torch.scatter(input: Tensor, dim: int, index: LongTensor, src: LongTensor)
    # * input, index and src tensors have same ndims.


        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))
        yield SampleInput(make_arg(shape, noncontiguous=True))

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

    if (
        enable_large_value_testing
        and dtype != torch.bool
        and dtype not in complex_dtypes
    ):
        filtered_large_values = _filter_domain(_large_values(dtype))
        yield SampleInput(
            torch.tensor(
                filtered_large_values,
                device="cuda",
                dtype=dtype,
                requires_grad=requires_grad,
            )
        )

    if enable_small_value_testing and dtype != torch.bool:
        filtered_small_values = _filter_domain(_small_values(dtype))
        yield SampleInput(
            torch.tensor(


    # Error: Trying to create tensor with negative dimension
    negative_input_shape = [2, -2]
    yield SampleInput(
        negative_input_shape, make_number(dtype), dtype
    ), RuntimeError, "The value -2 at index 1 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)."

def scatter_generator(
    op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs
):
    # torch.scatter(input: Tensor, dim: int, index: LongTensor, src: LongTensor)
    # * input, index and src tensors have same ndims.
    # * index tensors must be <= input tensor along all dims.
    # * index tensors must be == src tensor along all dims.
    # * index tensors must have unique value across specified axis.

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

    def make_unique_index(shape_b, dim, extent):
        logits_shape = list(shape_b)
        logits_shape[dim] = extent
        logits = make_tensor(logits_shape, device="cuda", dtype=torch.float)
        # return index tensor with unique entry
        return logits.argsort(dim).narrow(dim, 0, shape_b[dim])

    make_index = partial(
        make_tensor, device="cuda", dtype=torch.long, requires_grad=False
    )

    # a.shape, dim, b.shape
    cases = (
        ((8, 2, 3), 0, (8, 2, 3)),
        ((8, 2, 3), 1, (8, 2, 3)),
        ((8, 2, 3), 2, (8, 2, 3)),
        # TODO: enable the test below when we fix mapping for scatter.
        # scatter supporting unmatched scatter dim
        # ((8, 2, 3), 0, (4, 2, 3)),
        # ((8, 2, 3), 1, (8, 1, 3)),
        # ((8, 2, 3), 2, (8, 2, 2)),
        # ((8,), 0, (8)),
        # ((8,), 0, (4)),
        # ((8,), 0, (1)),
        # ((8, 2, 3), -3, (4, 2, 3)),
        # ((8, 2, 3), -2, (8, 1, 3)),


    # TODO: enable mxfp8 test when backend supports it.
    make_arg = partial(
        make_tensor,
        dtype=dtype,
        device="cuda",
        low=None,
        high=None,
        requires_grad=requires_grad,
    )

    make_scale_factor = partial(
        make_tensor,
        dtype=torch.float32,
        device="cuda",
        low=None,
        high=None,
        requires_grad=False,
    )

    def make_index(extent, num_groups):
        group_size = extent // num_groups
        return torch.arange(
            group_size,
            group_size * g + 1,
            group_size,
            device="cuda",
            dtype=torch.int32,
            requires_grad=False,
        )

    # TODO: expand the test when fallback kernel restrictions are lifted
    #       currently only bf16 output is supported.
    #       there are also restrictions on the input/output shapes.
    # Test various group sizes and matrix dimensions
    # configs: list(g, m, k, n, output_dtype)
    configs = (
        (4, 128, 256, 64, torch.bfloat16),
        (2, 32, 32, 32, torch.bfloat16),
    )

    # TODO: Enable mxfp8 test when backend supports it.
    for config in configs:
        g, m, k, n, dtype = config
        # case 1: 2d x 2d
        mat1 = make_arg((m, k))
        mat2 = make_arg((k, n))
        scale1 = make_scale_factor((g, m, 1))
        scale2 = make_scale_factor((g, 1, n))
        offsets = make_index(k, g)


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

# import nvfuser_direct first to conditionally avoid importing nvfuser
import nvfuser_direct  # noqa: F401,F403

import torch
import pytest
import numpy as np
from copy import deepcopy
from typing import Callable

from opinfo_fusion_definitions import default_fd_fn
from opinfo_framework import create_op_test
from opinfo_core import ReferenceType, OpInfo, SampleInput
from opinfos import opinfos
from opinfo_utils import (
    ArgumentType,
    is_tensor,
    requiresJAX,
)

import io
from contextlib import redirect_stdout, redirect_stderr
from nvfuser_direct import FusionDefinition
from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test
from python.direct_utils import check_captured_python_definition

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

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

    assert len(symbolic_parameter_list) >= len(args)

    result = []
    for arg_type, a in zip(symbolic_parameter_list, args):
        if arg_type == ArgumentType.Symbolic:
            if isinstance(a, list) and all(map(is_tensor, a)):
                result.extend(a)
            else:
                result.append(a)
    return result

# ****** Check an Operation's Results are Correct ******

def torch_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # TODO If dtype is fp16 or bf16, skip dtype check because nvfuser promotes
    # to fp32 but does not return original dtype.
    # TODO Add specific dtype tolerances
    torch.testing.assert_close(
        nvfuser_result, torch_result, equal_nan=True, atol=1e-3, rtol=0
    )

@requiresJAX
def jax_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    jax_sample = sample.jax()
    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")
    else:
        jax_result = torch.asarray(np_array, device="cuda")

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # NOTE: dtype is not checked because jax will translate int64, float64, and complex128 to int32, float32 and complex64
    torch.testing.assert_close(
        nvfuser_result, jax_result, equal_nan=True, atol=1e-3, rtol=0, check_dtype=False
    )

def python_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    # python reference function does not accept keyword arguments
    assert len(sample.kwargs) == 0

    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    # expect only single result from function
    assert len(nvfuser_result) == 1

    # convert tensor arguments into flat, python lists
    python_sample = sample.python()

    # apply reference to python lists
    python_result = map(nvf_op.reference, *python_sample.args)

    # create pytorch tensor
    np_array = np.array(list(python_result))
    if np_array.shape == ():
        python_result = torch.tensor(
            np_array.item(), dtype=nvfuser_result[0].dtype, device="cuda"
        )
    else:
        python_result = torch.asarray(
            np_array, dtype=nvfuser_result[0].dtype, device="cuda"
        )

    # reshape flat output tensor into expected shape
    torch.testing.assert_close(
        nvfuser_result[0],
        python_result.reshape(nvfuser_result[0].shape),
        equal_nan=True,
        atol=1e-3,
        rtol=0,
    )

def correctness_test_fn(
    reference_type: ReferenceType,
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_correctness_fn
        if nvf_op.fd_correctness_fn is not None
        else default_fd_fn
    )
    if reference_type == ReferenceType.Pytorch:
        return torch_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Jax:
        return jax_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Python:
        return python_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Numpy:
        pytest.xfail("Numpy feference functions are not supported.")
    else:
        pytest.xfail("Reference function is not defined for this correctness test.")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.sample_input_generator is not None and op.supports_direct_bindings
    )
)
def test_correctness(op: OpInfo, dtype: torch.dtype):
    for sample in op.sample_input_generator(op, dtype):
        result = correctness_test_fn(op.reference_type, op, sample)
        if result is not None:
            return result

# ****** Check that an Operation's API Gives Appropriate Input Errors ******

def errors_test_fn(
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_error_input_fn
        if nvf_op.fd_error_input_fn is not None
        else default_fd_fn
    )
    with FusionDefinition() as fd:
        _fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

# A pair of parentheses ()/[] represents a capture group in regex.
# Escape parenthesis in regex string to match raw characters.
def _regex_escape_parenthesis(a: str) -> str:
    b = a.replace(r"[", r"\[").replace(r"]", r"\]")
    return b.replace(r"(", r"\(").replace(r")", r"\)")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.error_input_generator is not None and op.supports_direct_bindings
    )
)
@retry_on_oom_or_skip_test
def test_errors(op: OpInfo, dtype: torch.dtype):
    for sample, exception_type, exception_regex in op.error_input_generator(op, dtype):
        # TODO skip regex check because direct bindings do not have same error message as frontend.
        # Capture error strings to avoid false positives in the CI
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        with pytest.raises(exception_type), redirect_stdout(
            stdout_capture
        ), redirect_stderr(stderr_capture):
            errors_test_fn(op, sample)


from copy import deepcopy
from typing import Callable

from opinfo_fusion_definitions import default_fd_fn
from opinfo_framework import create_op_test
from opinfo_core import ReferenceType, OpInfo, SampleInput
from opinfos import opinfos
from opinfo_utils import (
    ArgumentType,
    is_tensor,
    requiresJAX,
)

import io
from contextlib import redirect_stdout, redirect_stderr
from nvfuser_direct import FusionDefinition
from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test
from python.direct_utils import check_captured_python_definition

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

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

    assert len(symbolic_parameter_list) >= len(args)

    result = []
    for arg_type, a in zip(symbolic_parameter_list, args):
        if arg_type == ArgumentType.Symbolic:
            if isinstance(a, list) and all(map(is_tensor, a)):
                result.extend(a)
            else:
                result.append(a)
    return result

# ****** Check an Operation's Results are Correct ******

def torch_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]


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

    assert len(symbolic_parameter_list) >= len(args)

    result = []
    for arg_type, a in zip(symbolic_parameter_list, args):
        if arg_type == ArgumentType.Symbolic:
            if isinstance(a, list) and all(map(is_tensor, a)):
                result.extend(a)
            else:
                result.append(a)
    return result

# ****** Check an Operation's Results are Correct ******

def torch_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # TODO If dtype is fp16 or bf16, skip dtype check because nvfuser promotes
    # to fp32 but does not return original dtype.
    # TODO Add specific dtype tolerances
    torch.testing.assert_close(
        nvfuser_result, torch_result, equal_nan=True, atol=1e-3, rtol=0
    )

@requiresJAX
def jax_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    jax_sample = sample.jax()
    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")


    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # TODO If dtype is fp16 or bf16, skip dtype check because nvfuser promotes
    # to fp32 but does not return original dtype.
    # TODO Add specific dtype tolerances
    torch.testing.assert_close(
        nvfuser_result, torch_result, equal_nan=True, atol=1e-3, rtol=0
    )

@requiresJAX
def jax_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    jax_sample = sample.jax()
    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")
    else:
        jax_result = torch.asarray(np_array, device="cuda")

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # NOTE: dtype is not checked because jax will translate int64, float64, and complex128 to int32, float32 and complex64
    torch.testing.assert_close(
        nvfuser_result, jax_result, equal_nan=True, atol=1e-3, rtol=0, check_dtype=False
    )

def python_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    # python reference function does not accept keyword arguments
    assert len(sample.kwargs) == 0

    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    # expect only single result from function
    assert len(nvfuser_result) == 1

    # convert tensor arguments into flat, python lists
    python_sample = sample.python()

    # apply reference to python lists


    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")
    else:
        jax_result = torch.asarray(np_array, device="cuda")

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # NOTE: dtype is not checked because jax will translate int64, float64, and complex128 to int32, float32 and complex64
    torch.testing.assert_close(
        nvfuser_result, jax_result, equal_nan=True, atol=1e-3, rtol=0, check_dtype=False
    )

def python_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    # python reference function does not accept keyword arguments
    assert len(sample.kwargs) == 0

    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args)
    inputs = parse_args_fusion_execution(nvf_op, *sample.args)
    inputs_cap = deepcopy(inputs)
    nvfuser_result = fd.execute(inputs)
    assert check_captured_python_definition(nvfuser_result, fd, inputs_cap)

    # expect only single result from function
    assert len(nvfuser_result) == 1

    # convert tensor arguments into flat, python lists
    python_sample = sample.python()

    # apply reference to python lists
    python_result = map(nvf_op.reference, *python_sample.args)

    # create pytorch tensor
    np_array = np.array(list(python_result))
    if np_array.shape == ():
        python_result = torch.tensor(
            np_array.item(), dtype=nvfuser_result[0].dtype, device="cuda"
        )
    else:
        python_result = torch.asarray(
            np_array, dtype=nvfuser_result[0].dtype, device="cuda"
        )

    # reshape flat output tensor into expected shape
    torch.testing.assert_close(
        nvfuser_result[0],
        python_result.reshape(nvfuser_result[0].shape),
        equal_nan=True,
        atol=1e-3,
        rtol=0,
    )

def correctness_test_fn(
    reference_type: ReferenceType,
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_correctness_fn
        if nvf_op.fd_correctness_fn is not None
        else default_fd_fn
    )
    if reference_type == ReferenceType.Pytorch:
        return torch_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Jax:
        return jax_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Python:
        return python_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Numpy:
        pytest.xfail("Numpy feference functions are not supported.")


    np_array = np.array(list(python_result))
    if np_array.shape == ():
        python_result = torch.tensor(
            np_array.item(), dtype=nvfuser_result[0].dtype, device="cuda"
        )
    else:
        python_result = torch.asarray(
            np_array, dtype=nvfuser_result[0].dtype, device="cuda"
        )

    # reshape flat output tensor into expected shape
    torch.testing.assert_close(
        nvfuser_result[0],
        python_result.reshape(nvfuser_result[0].shape),
        equal_nan=True,
        atol=1e-3,
        rtol=0,
    )

def correctness_test_fn(
    reference_type: ReferenceType,
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_correctness_fn
        if nvf_op.fd_correctness_fn is not None
        else default_fd_fn
    )
    if reference_type == ReferenceType.Pytorch:
        return torch_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Jax:
        return jax_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Python:
        return python_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Numpy:
        pytest.xfail("Numpy feference functions are not supported.")
    else:
        pytest.xfail("Reference function is not defined for this correctness test.")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.sample_input_generator is not None and op.supports_direct_bindings
    )
)
def test_correctness(op: OpInfo, dtype: torch.dtype):
    for sample in op.sample_input_generator(op, dtype):
        result = correctness_test_fn(op.reference_type, op, sample)
        if result is not None:
            return result

# ****** Check that an Operation's API Gives Appropriate Input Errors ******

def errors_test_fn(


    )
    if reference_type == ReferenceType.Pytorch:
        return torch_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Jax:
        return jax_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Python:
        return python_correctness_test_fn(_fd_fn, nvf_op, sample)
    elif reference_type == ReferenceType.Numpy:
        pytest.xfail("Numpy feference functions are not supported.")
    else:
        pytest.xfail("Reference function is not defined for this correctness test.")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.sample_input_generator is not None and op.supports_direct_bindings
    )
)
def test_correctness(op: OpInfo, dtype: torch.dtype):
    for sample in op.sample_input_generator(op, dtype):
        result = correctness_test_fn(op.reference_type, op, sample)
        if result is not None:
            return result

# ****** Check that an Operation's API Gives Appropriate Input Errors ******

def errors_test_fn(
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_error_input_fn
        if nvf_op.fd_error_input_fn is not None
        else default_fd_fn
    )
    with FusionDefinition() as fd:
        _fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

# A pair of parentheses ()/[] represents a capture group in regex.


        pytest.xfail("Reference function is not defined for this correctness test.")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.sample_input_generator is not None and op.supports_direct_bindings
    )
)
def test_correctness(op: OpInfo, dtype: torch.dtype):
    for sample in op.sample_input_generator(op, dtype):
        result = correctness_test_fn(op.reference_type, op, sample)
        if result is not None:
            return result

# ****** Check that an Operation's API Gives Appropriate Input Errors ******

def errors_test_fn(
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_error_input_fn
        if nvf_op.fd_error_input_fn is not None
        else default_fd_fn
    )
    with FusionDefinition() as fd:
        _fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

# A pair of parentheses ()/[] represents a capture group in regex.
# Escape parenthesis in regex string to match raw characters.
def _regex_escape_parenthesis(a: str) -> str:
    b = a.replace(r"[", r"\[").replace(r"]", r"\]")
    return b.replace(r"(", r"\(").replace(r")", r"\)")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.error_input_generator is not None and op.supports_direct_bindings
    )
)
@retry_on_oom_or_skip_test
def test_errors(op: OpInfo, dtype: torch.dtype):
    for sample, exception_type, exception_regex in op.error_input_generator(op, dtype):
        # TODO skip regex check because direct bindings do not have same error message as frontend.


# ****** Check that an Operation's API Gives Appropriate Input Errors ******

def errors_test_fn(
    nvf_op: OpInfo,
    sample: SampleInput,
):
    _fd_fn = (
        nvf_op.fd_error_input_fn
        if nvf_op.fd_error_input_fn is not None
        else default_fd_fn
    )
    with FusionDefinition() as fd:
        _fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

# A pair of parentheses ()/[] represents a capture group in regex.
# Escape parenthesis in regex string to match raw characters.
def _regex_escape_parenthesis(a: str) -> str:
    b = a.replace(r"[", r"\[").replace(r"]", r"\]")
    return b.replace(r"(", r"\(").replace(r")", r"\)")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.error_input_generator is not None and op.supports_direct_bindings
    )
)
@retry_on_oom_or_skip_test
def test_errors(op: OpInfo, dtype: torch.dtype):
    for sample, exception_type, exception_regex in op.error_input_generator(op, dtype):
        # TODO skip regex check because direct bindings do not have same error message as frontend.
        # Capture error strings to avoid false positives in the CI
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        with pytest.raises(exception_type), redirect_stdout(
            stdout_capture
        ), redirect_stderr(stderr_capture):
            errors_test_fn(op, sample)


    with FusionDefinition() as fd:
        _fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

# A pair of parentheses ()/[] represents a capture group in regex.
# Escape parenthesis in regex string to match raw characters.
def _regex_escape_parenthesis(a: str) -> str:
    b = a.replace(r"[", r"\[").replace(r"]", r"\]")
    return b.replace(r"(", r"\(").replace(r")", r"\)")

@create_op_test(
    tuple(
        op
        for op in opinfos
        if op.error_input_generator is not None and op.supports_direct_bindings
    )
)
@retry_on_oom_or_skip_test
def test_errors(op: OpInfo, dtype: torch.dtype):
    for sample, exception_type, exception_regex in op.error_input_generator(op, dtype):
        # TODO skip regex check because direct bindings do not have same error message as frontend.
        # Capture error strings to avoid false positives in the CI
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        with pytest.raises(exception_type), redirect_stdout(
            stdout_capture
        ), redirect_stderr(stderr_capture):
            errors_test_fn(op, sample)


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

import inspect
import torch
from typing import Callable
from opinfo_utils import map_dtype_to_str
import pytest

def _instantiate_opinfo_test_template(
    template: Callable, *, opinfo, dtype: torch.dtype
) -> Callable:
    """Instantiates a test template for an operator."""

    def test():
        # Ref: https://github.com/pytorch/pytorch/blob/aa8ea1d787a9d21b064b664c5344376265feea6c/torch/testing/_internal/common_utils.py#L2251-L2263
        # > CUDA device side error will cause subsequence test cases to fail.
        # > stop entire test suite if catches RuntimeError during torch.cuda.synchronize().
        if torch.cuda.is_initialized():
            try:
                torch.cuda.synchronize()
            except RuntimeError as rte:
                pytest.exit(
                    "TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure"
                )

        return template(opinfo, dtype)

    test.__name__ = "_".join((template.__name__, opinfo.name, map_dtype_to_str[dtype]))
    test.__module__ = test.__module__
    return test

class create_op_test:
    def __init__(self, opinfos, *, scope=None):
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes.
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)
        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(
                    test_template,
                    opinfo=opinfo,
                    dtype=dtype,
                )
                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test

# This pseudo-decorator enables automatic serialization upon program exit and
# tests deserializing the default workspace upon creating the tests.
#
# Serializing error test cases corrupts the serialized binary. We call
# FusionCache.reset() to clear the cache after running an error test in
# `test_python_frontend.py'. In the pytest framework, the error tests are
# separate from the correctness tests. Only apply this decorator to the
# correctness tests to avoid calling FusionCache.reset().
class atexit_serde_create_op_test(create_op_test):
    def __init__(self, opinfos, *, scope=None):
        from python.utils import atexit_serde_check

        atexit_serde_check()

        # Replicate create_op_test.__init__ to have correct scope
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope


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

import inspect
import torch
from typing import Callable
from opinfo_utils import map_dtype_to_str
import pytest

def _instantiate_opinfo_test_template(
    template: Callable, *, opinfo, dtype: torch.dtype
) -> Callable:
    """Instantiates a test template for an operator."""

    def test():
        # Ref: https://github.com/pytorch/pytorch/blob/aa8ea1d787a9d21b064b664c5344376265feea6c/torch/testing/_internal/common_utils.py#L2251-L2263
        # > CUDA device side error will cause subsequence test cases to fail.
        # > stop entire test suite if catches RuntimeError during torch.cuda.synchronize().
        if torch.cuda.is_initialized():
            try:
                torch.cuda.synchronize()
            except RuntimeError as rte:
                pytest.exit(
                    "TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure"
                )

        return template(opinfo, dtype)

    test.__name__ = "_".join((template.__name__, opinfo.name, map_dtype_to_str[dtype]))
    test.__module__ = test.__module__
    return test

class create_op_test:
    def __init__(self, opinfos, *, scope=None):
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes.
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)
        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(


    def test():
        # Ref: https://github.com/pytorch/pytorch/blob/aa8ea1d787a9d21b064b664c5344376265feea6c/torch/testing/_internal/common_utils.py#L2251-L2263
        # > CUDA device side error will cause subsequence test cases to fail.
        # > stop entire test suite if catches RuntimeError during torch.cuda.synchronize().
        if torch.cuda.is_initialized():
            try:
                torch.cuda.synchronize()
            except RuntimeError as rte:
                pytest.exit(
                    "TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure"
                )

        return template(opinfo, dtype)

    test.__name__ = "_".join((template.__name__, opinfo.name, map_dtype_to_str[dtype]))
    test.__module__ = test.__module__
    return test

class create_op_test:
    def __init__(self, opinfos, *, scope=None):
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes.
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)
        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(
                    test_template,
                    opinfo=opinfo,
                    dtype=dtype,
                )
                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test

# This pseudo-decorator enables automatic serialization upon program exit and
# tests deserializing the default workspace upon creating the tests.
#
# Serializing error test cases corrupts the serialized binary. We call
# FusionCache.reset() to clear the cache after running an error test in
# `test_python_frontend.py'. In the pytest framework, the error tests are
# separate from the correctness tests. Only apply this decorator to the
# correctness tests to avoid calling FusionCache.reset().
class atexit_serde_create_op_test(create_op_test):
    def __init__(self, opinfos, *, scope=None):
        from python.utils import atexit_serde_check

        atexit_serde_check()

        # Replicate create_op_test.__init__ to have correct scope
        self.opinfos = opinfos

        # Acquires the caller's global scope


        #   functions are directly assigned to the requested scope (the caller's global scope by default)
        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(
                    test_template,
                    opinfo=opinfo,
                    dtype=dtype,
                )
                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test

# This pseudo-decorator enables automatic serialization upon program exit and
# tests deserializing the default workspace upon creating the tests.
#
# Serializing error test cases corrupts the serialized binary. We call
# FusionCache.reset() to clear the cache after running an error test in
# `test_python_frontend.py'. In the pytest framework, the error tests are
# separate from the correctness tests. Only apply this decorator to the
# correctness tests to avoid calling FusionCache.reset().
class atexit_serde_create_op_test(create_op_test):
    def __init__(self, opinfos, *, scope=None):
        from python.utils import atexit_serde_check

        atexit_serde_check()

        # Replicate create_op_test.__init__ to have correct scope
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope


test

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

import inspect
import torch
from typing import Callable
from opinfo_utils import map_dtype_to_str
import pytest

def _instantiate_opinfo_test_template(
    template: Callable, *, opinfo, dtype: torch.dtype
) -> Callable:
    """Instantiates a test template for an operator."""

    def test():
        # Ref: https://github.com/pytorch/pytorch/blob/aa8ea1d787a9d21b064b664c5344376265feea6c/torch/testing/_internal/common_utils.py#L2251-L2263
        # > CUDA device side error will cause subsequence test cases to fail.
        # > stop entire test suite if catches RuntimeError during torch.cuda.synchronize().
        if torch.cuda.is_initialized():
            try:
                torch.cuda.synchronize()
            except RuntimeError as rte:
                pytest.exit(
                    "TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure"
                )

        return template(opinfo, dtype)

    test.__name__ = "_".join((template.__name__, opinfo.name, map_dtype_to_str[dtype]))
    test.__module__ = test.__module__
    return test

class create_op_test:
    def __init__(self, opinfos, *, scope=None):
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes.
        #   Since Python doesn't natively support one-to-many function decorators, the produced


        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(
                    test_template,
                    opinfo=opinfo,
                    dtype=dtype,
                )
                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test

# This pseudo-decorator enables automatic serialization upon program exit and
# tests deserializing the default workspace upon creating the tests.
#
# Serializing error test cases corrupts the serialized binary. We call
# FusionCache.reset() to clear the cache after running an error test in
# `test_python_frontend.py'. In the pytest framework, the error tests are
# separate from the correctness tests. Only apply this decorator to the
# correctness tests to avoid calling FusionCache.reset().
class atexit_serde_create_op_test(create_op_test):
    def __init__(self, opinfos, *, scope=None):
        from python.utils import atexit_serde_check

        atexit_serde_check()

        # Replicate create_op_test.__init__ to have correct scope
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope


                    "TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure"
                )

        return template(opinfo, dtype)

    test.__name__ = "_".join((template.__name__, opinfo.name, map_dtype_to_str[dtype]))
    test.__module__ = test.__module__
    return test

class create_op_test:
    def __init__(self, opinfos, *, scope=None):
        self.opinfos = opinfos

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes.
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)
        for opinfo in self.opinfos:
            for dtype in sorted(opinfo.dtypes, key=lambda t: repr(t)):
                test = _instantiate_opinfo_test_template(
                    test_template,
                    opinfo=opinfo,
                    dtype=dtype,
                )
                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test

# This pseudo-decorator enables automatic serialization upon program exit and
# tests deserializing the default workspace upon creating the tests.
#
# Serializing error test cases corrupts the serialized binary. We call
# FusionCache.reset() to clear the cache after running an error test in
# `test_python_frontend.py'. In the pytest framework, the error tests are
# separate from the correctness tests. Only apply this decorator to the
# correctness tests to avoid calling FusionCache.reset().
class atexit_serde_create_op_test(create_op_test):
    def __init__(self, opinfos, *, scope=None):
        from python.utils import atexit_serde_check

        atexit_serde_check()

        # Replicate create_op_test.__init__ to have correct scope
        self.opinfos = opinfos

        # Acquires the caller's global scope


tests.python.opinfo.opinfos

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

import math
import torch
from looseversion import LooseVersion
from opinfo_core import OpInfo, ReferenceType, Domain
from opinfo_fusion_definitions import (
    api_test_fd_fn,
    tensor_input_fd_fn,
    tensor_api_test_fd_fn,
    vector_api_test_fd_fn,
)
from opinfo_input_generators import (
    argsort_generator,
    topk_generator,
    topk_error_generator,
    broadcast_error_generator,
    broadcast_in_dim_generator,
    broadcast_in_dim_error_generator,
    cat_generator,
    cat_error_generator,
    div_input_generator,
    define_tensor_error_generator,
    define_vector_constant_error_generator,
    elementwise_binary_generator,
    _elementwise_binary_torch,
    elementwise_unary_generator,
    _elementwise_unary_torch,
    full_error_generator,
    gather_generator,
    scatter_generator,
    index_select_generator,
    index_select_error_generator,
    index_put_accumulate_generator,
    iota_error_generator,
    pad_error_generator,
    permute_generator,
    permute_error_generator,
    random_dist_error_generator,
    reduction_error_generator,
    reshape_generator,
    reshape_error_generator,
    slice_generator,
    slice_error_generator,
    squeeze_generator,
    squeeze_error_generator,
    take_along_axis_generator,
    take_along_axis_error_generator,
    tensor_size_error_generator,
    var_mean_generator,
    vector_at_error_generator,
    where_error_generator,
    matmul_input_generator,
    linear_input_generator,
    linear_error_generator,
    triu_input_generator,
    triu_error_generator,
    cumsum_input_generator,
    cumsum_error_generator,
    grouped_mm_input_generator,
    scaled_grouped_mm_input_generator,
    scaled_mm_input_generator,
)
from opinfo_utils import (
    bool_int_dtypes,
    complex_dtypes,
    full_precision_float_dtypes,
    int_dtypes,
    int_float_dtypes,
    float_dtypes,
    float_complex_dtypes,
    ArgumentType,
    JAX_AVAILABLE,
)
from functools import partial

if JAX_AVAILABLE:
    import jax

eps = 1e-2

opinfos = []

""" Start Fusion Input Operations """
fusion_input_ops = []

define_tensor_opinfo = OpInfo(
    lambda fd: fd.define_tensor,
    "define_tensor",
    error_input_generator=define_tensor_error_generator,
    fd_error_input_fn=tensor_input_fd_fn,
)
fusion_input_ops.append(define_tensor_opinfo)

# NOTE: "define_vector" only supports vectors of integers that represent
# tensor shapes and is not a general interface for defining vectors of
# data.  Vectors of data should be handled with a 1D `define_tensor`.
define_vector_constant_opinfo = OpInfo(
    lambda fd: fd.define_vector,
    "define_vector_constant",
    error_input_generator=define_vector_constant_error_generator,
    fd_error_input_fn=api_test_fd_fn,
    # Direct bindings supports python lists and tuples of values.
    # These python lists are directly indexable, so `define_vector_constant` is not needed.
    supports_direct_bindings=False,
)
fusion_input_ops.append(define_vector_constant_opinfo)

""" End Fusion Input Operations """

""" Start Unary-Float Operations """
unary_ops = []

abs_opinfo = OpInfo(
    lambda fd: fd.ops.abs,
    "abs",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.abs),
    is_clonable=True,
)
unary_ops.append(abs_opinfo)

acos_opinfo = OpInfo(
    lambda fd: fd.ops.acos,
    "acos",
    domain=Domain(-1, 1),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.acos),
    is_clonable=True,
)
unary_ops.append(acos_opinfo)

acosh_opinfo = OpInfo(
    lambda fd: fd.ops.acosh,
    "acosh",
    domain=Domain(-1, math.inf),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.acosh),
    is_clonable=True,
)
unary_ops.append(acosh_opinfo)

asin_opinfo = OpInfo(
    lambda fd: fd.ops.asin,
    "asin",
    domain=Domain(-1, 1),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.asin),
    is_clonable=True,
)
unary_ops.append(asin_opinfo)

asinh_opinfo = OpInfo(
    lambda fd: fd.ops.asinh,
    "asinh",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.asinh),
    is_clonable=True,
)
unary_ops.append(asinh_opinfo)

atan_opinfo = OpInfo(
    lambda fd: fd.ops.atan,
    "atan",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.atan),
    is_clonable=True,
)
unary_ops.append(atan_opinfo)

atanh_opinfo = OpInfo(
    lambda fd: fd.ops.atanh,
    "atanh",
    domain=Domain(-1 + eps, 1 + eps),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.atanh),
    is_clonable=True,
)
unary_ops.append(atanh_opinfo)

bitwise_not_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_not,
    "bitwise_not",
    dtypes=bool_int_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.bitwise_not),
    is_clonable=True,
)
unary_ops.append(bitwise_not_opinfo)

# TODO add nvfuser exception for int dtypes
ceil_opinfo = OpInfo(
    lambda fd: fd.ops.ceil,
    "ceil",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.ceil),
    is_clonable=True,
)
unary_ops.append(ceil_opinfo)

cos_opinfo = OpInfo(
    lambda fd: fd.ops.cos,
    "cos",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.cos),
    is_clonable=True,
)
unary_ops.append(cos_opinfo)

cosh_opinfo = OpInfo(
    lambda fd: fd.ops.cosh,
    "cosh",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.cosh),
    is_clonable=True,
)
unary_ops.append(cosh_opinfo)

erf_opinfo = OpInfo(
    lambda fd: fd.ops.erf,
    "erf",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.erf),
    is_clonable=True,
)
unary_ops.append(erf_opinfo)

erfc_opinfo = OpInfo(
    lambda fd: fd.ops.erfc,
    "erfc",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.erfc),
    is_clonable=True,
)
unary_ops.append(erfc_opinfo)

erfcinv_opinfo = OpInfo(
    lambda fd: fd.ops.erfcinv,
    "erfcinv",
    dtypes=(
        torch.float32,
        torch.float64,
    ),
    domain=Domain(0.3, 0.7),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(lambda x: torch.erfinv(1 - x)),
    is_clonable=True,
)
unary_ops.append(erfcinv_opinfo)

erfinv_opinfo = OpInfo(
    lambda fd: fd.ops.erfinv,
    "erfinv",
    dtypes=int_float_dtypes,
    domain=Domain(-1, 1),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.erfinv),
    is_clonable=True,
)
unary_ops.append(erfinv_opinfo)

exp_opinfo = OpInfo(
    lambda fd: fd.ops.exp,
    "exp",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.exp),
    is_clonable=True,
)
unary_ops.append(exp_opinfo)

exp2_opinfo = OpInfo(
    lambda fd: fd.ops.exp2,
    "exp2",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.exp2),
    is_clonable=True,
)
unary_ops.append(exp2_opinfo)

expm1_opinfo = OpInfo(
    lambda fd: fd.ops.expm1,
    "expm1",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.expm1),
    is_clonable=True,
)
unary_ops.append(expm1_opinfo)

# TODO add nvfuser exception for int dtypes
floor_opinfo = OpInfo(
    lambda fd: fd.ops.floor,
    "floor",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.floor),
    is_clonable=True,
)
unary_ops.append(floor_opinfo)

frac_opinfo = OpInfo(
    lambda fd: fd.ops.frac,
    "frac",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.frac),
    is_clonable=True,
)
unary_ops.append(frac_opinfo)

isfinite_opinfo = OpInfo(
    lambda fd: fd.ops.isfinite,
    "isfinite",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isfinite),
    is_clonable=True,
)
unary_ops.append(isfinite_opinfo)

isinf_opinfo = OpInfo(
    lambda fd: fd.ops.isinf,
    "isinf",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isinf),
    is_clonable=True,
)
unary_ops.append(isinf_opinfo)

isnan_opinfo = OpInfo(
    lambda fd: fd.ops.isnan,
    "isnan",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isnan),
    is_clonable=True,
)
unary_ops.append(isnan_opinfo)

# NOTE half-precision floating types are not automatically promoted to fp32
isneginf_opinfo = OpInfo(
    lambda fd: fd.ops.isneginf,
    "isneginf",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isneginf),
    is_clonable=True,
)
unary_ops.append(isneginf_opinfo)

# NOTE half-precision floating types are not automatically promoted to fp32
isposinf_opinfo = OpInfo(
    lambda fd: fd.ops.isposinf,
    "isposinf",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isposinf),
    is_clonable=True,
)
unary_ops.append(isposinf_opinfo)

isreal_opinfo = OpInfo(
    lambda fd: fd.ops.isreal,
    "isreal",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.isreal),
    is_clonable=True,
)
unary_ops.append(isreal_opinfo)

lgamma_opinfo = OpInfo(
    lambda fd: fd.ops.lgamma,
    "lgamma",
    dtypes=int_float_dtypes,
    domain=Domain(-1.0 + eps, math.inf),
    sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
    reference=_elementwise_unary_torch(torch.lgamma),
    is_clonable=True,
)
unary_ops.append(lgamma_opinfo)

log_opinfo = OpInfo(
    lambda fd: fd.ops.log,
    "log",
    domain=Domain(0, math.inf),
    sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
    reference=_elementwise_unary_torch(torch.log),
    is_clonable=True,
)
unary_ops.append(log_opinfo)

log10_opinfo = OpInfo(
    lambda fd: fd.ops.log10,
    "log10",
    dtypes=int_float_dtypes,
    domain=Domain(0, math.inf),
    sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
    reference=_elementwise_unary_torch(torch.log10),
    is_clonable=True,
)
unary_ops.append(log10_opinfo)

log1p_opinfo = OpInfo(
    lambda fd: fd.ops.log1p,
    "log1p",
    dtypes=int_float_dtypes,
    domain=Domain(-1 + eps, math.inf),
    sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
    reference=_elementwise_unary_torch(torch.log1p),
    is_clonable=True,
)
unary_ops.append(log1p_opinfo)

log2_opinfo = OpInfo(
    lambda fd: fd.ops.log2,
    "log2",
    domain=Domain(0, math.inf),
    sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
    reference=_elementwise_unary_torch(torch.log2),
    is_clonable=True,
)
unary_ops.append(log2_opinfo)

neg_opinfo = OpInfo(
    lambda fd: fd.ops.neg,
    "neg",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.neg),
    is_clonable=True,
)
unary_ops.append(neg_opinfo)

reciprocal_opinfo = OpInfo(
    lambda fd: fd.ops.reciprocal,
    "reciprocal",
    domain=Domain(0 + eps, math.inf),
    sample_input_generator=partial(
        elementwise_unary_generator,
        enable_small_value_testing=False,
        enable_extremal_value_testing=False,
        exclude_zero=True,
    ),
    reference=_elementwise_unary_torch(torch.reciprocal),
    is_clonable=True,
)
unary_ops.append(reciprocal_opinfo)

# TODO add nvfuser exception for int dtypes
round_opinfo = OpInfo(
    lambda fd: fd.ops.round,
    "round",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.round),
    is_clonable=True,
)
unary_ops.append(round_opinfo)

rsqrt_opinfo = OpInfo(
    lambda fd: fd.ops.rsqrt,
    "rsqrt",
    domain=Domain(0 + eps, math.inf),
    sample_input_generator=partial(
        elementwise_unary_generator,
        enable_small_value_testing=False,
        enable_extremal_value_testing=False,
        exclude_zero=True,
    ),
    reference=_elementwise_unary_torch(torch.rsqrt),
    is_clonable=True,
)
unary_ops.append(rsqrt_opinfo)

sigmoid_opinfo = OpInfo(
    lambda fd: fd.ops.sigmoid,
    "sigmoid",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.sigmoid),
    is_clonable=True,
)
unary_ops.append(sigmoid_opinfo)

signbit_opinfo = OpInfo(
    lambda fd: fd.ops.signbit,
    "signbit",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.signbit),
    is_clonable=True,
)
unary_ops.append(signbit_opinfo)

sin_opinfo = OpInfo(
    lambda fd: fd.ops.sin,
    "sin",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.sin),
    is_clonable=True,
)
unary_ops.append(sin_opinfo)

sinh_opinfo = OpInfo(
    lambda fd: fd.ops.sinh,
    "sinh",
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.sinh),
    is_clonable=True,
)
unary_ops.append(sinh_opinfo)

sqrt_opinfo = OpInfo(
    lambda fd: fd.ops.sqrt,
    "sqrt",
    domain=Domain(0, math.inf),
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.sqrt),
    is_clonable=True,
)
unary_ops.append(sqrt_opinfo)

tan_opinfo = OpInfo(
    lambda fd: fd.ops.tan,
    "tan",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.tan),
    is_clonable=True,
)
unary_ops.append(tan_opinfo)

tanh_opinfo = OpInfo(
    lambda fd: fd.ops.tanh,
    "tanh",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.tanh),
    is_clonable=True,
)
unary_ops.append(tanh_opinfo)

# TODO add nvfuser exception for int dtypes
trunc_opinfo = OpInfo(
    lambda fd: fd.ops.trunc,
    "trunc",
    dtypes=full_precision_float_dtypes,
    sample_input_generator=elementwise_unary_generator,
    reference=_elementwise_unary_torch(torch.trunc),
    is_clonable=True,
)
unary_ops.append(trunc_opinfo)

""" End Unary-Float Operations """

""" Start Binary Operations """

# atan2 --- promote int to float; allows fp16 and bf16
# nextafter, truediv  --- promote int to float; requires full-precision fp32, fp64
# ceildiv, div, fmod, mod, remainder, truediv --- except_zero
# add, mul, pow, sub
# bitwise_and, bitwise_or, bitwise_xor --- bool_int_only
# bitwise_left_shift, bitwise_right_shift, logical_right_shift --- int_only
# eq, ne, ge, gt, le, lt --- compare

# TODO Add "ceildiv" to python_frontend
# TODO Add support for python reference for "mod".
# TODO atan2 - complex dtypes are unsupported, but we fail when compiling kernel
# TODO logical_right_shift - domain of shift parameter is non-zero; Otherwise the result is undefined.

binary_ops = []

add_opinfo = OpInfo(
    lambda fd: fd.ops.add,
    "add",
    sample_input_generator=partial(
        elementwise_binary_generator, enable_extremal_value_testing=False
    ),
    reference=_elementwise_binary_torch(torch.add),
    is_clonable=True,
)
binary_ops.append(add_opinfo)

# TODO complex dtypes are unsupported, but we fail when compiling kernel
atan2_opinfo = OpInfo(
    lambda fd: fd.ops.atan2,
    "atan2",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.atan2),
    is_clonable=True,
)
binary_ops.append(atan2_opinfo)

bitwise_and_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_and,
    "bitwise_and",
    dtypes=bool_int_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.bitwise_and),
    is_clonable=True,
)
binary_ops.append(bitwise_and_opinfo)

bitwise_left_shift_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_left_shift,
    "bitwise_left_shift",
    dtypes=int_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.bitwise_left_shift),
    is_clonable=True,
)
binary_ops.append(bitwise_left_shift_opinfo)

bitwise_or_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_or,
    "bitwise_or",
    dtypes=bool_int_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.bitwise_or),
    is_clonable=True,
)
binary_ops.append(bitwise_or_opinfo)

bitwise_right_shift_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_right_shift,
    "bitwise_right_shift",
    dtypes=int_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.bitwise_right_shift),
    is_clonable=True,
)
binary_ops.append(bitwise_right_shift_opinfo)

bitwise_xor_opinfo = OpInfo(
    lambda fd: fd.ops.bitwise_xor,
    "bitwise_xor",
    dtypes=bool_int_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.bitwise_xor),
    is_clonable=True,
)
binary_ops.append(bitwise_xor_opinfo)

div_opinfo = OpInfo(
    lambda fd: fd.ops.div,
    "div",
    dtypes=float_complex_dtypes,
    sample_input_generator=div_input_generator,
    reference=_elementwise_binary_torch(torch.div),
    is_clonable=True,
)
binary_ops.append(div_opinfo)

eq_opinfo = OpInfo(
    lambda fd: fd.ops.eq,
    "eq",
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.eq),
    is_clonable=True,
)
binary_ops.append(eq_opinfo)

fmod_opinfo = OpInfo(
    lambda fd: fd.ops.fmod,
    "fmod",
    dtypes=int_float_dtypes,
    sample_input_generator=partial(elementwise_binary_generator, exclude_zero=True),
    reference=_elementwise_binary_torch(torch.fmod),
    is_clonable=True,
)
binary_ops.append(fmod_opinfo)

ge_opinfo = OpInfo(
    lambda fd: fd.ops.ge,
    "ge",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.ge),
    is_clonable=True,
)
binary_ops.append(ge_opinfo)

gt_opinfo = OpInfo(
    lambda fd: fd.ops.gt,
    "gt",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.gt),
    is_clonable=True,
)
binary_ops.append(gt_opinfo)

le_opinfo = OpInfo(
    lambda fd: fd.ops.le,
    "le",
    dtypes=int_float_dtypes,
    sample_input_generator=elementwise_binary_generator,
    reference=_elementwise_binary_torch(torch.le),
    is_clonable=True,
)
binary_ops.append(le_opinfo)

# TODO domain of shift parameter greater than zero; Otherwise the result is undefined.
logical_right_shift_opinfo = OpInfo(
    lambda fd: fd.ops.logical_right_shift,
    "logical_right_shift",
    domain=Domain(0, None),
    dtypes=int_dtypes,
    sample_input_generator=partial(
        elementwise_binary_generator,
        enable_broadcast_testing=False,
        enable_extremal_value_testing=False,
        enable_large_value_testing=False,
        enable_small_value_testing=False,
    ),
    reference=jax.lax.shift_right_logical if JAX_AVAILA
# ... truncated ...


broadcast_in_dim_constant_opinfo = OpInfo(
    lambda fd: fd.ops.broadcast_in_dim,
    "broadcast_in_dim_constant",
    sample_input_generator=broadcast_in_dim_generator,
    error_input_generator=broadcast_in_dim_error_generator,
    reference=jax.lax.broadcast_in_dim if JAX_AVAILABLE else None,
    reference_type=ReferenceType.Jax,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        # This argument is purposely Constant even though the positional
        # argument can also be symbolic.
        ArgumentType.Constant,
        ArgumentType.Constant,
    ),
)
shape_ops.append(broadcast_in_dim_constant_opinfo)

# NOTE: The symbolic version of broadcast_in_dim opinfo tests the "shape"
# argument with a Vector generated from another operation like ops.shape.
def broadcast_in_dim_sym_fn(fd, arg1, arg2, broadcast_dims):
    return fd.ops.broadcast_in_dim(arg1, arg2.shape(), broadcast_dims)

def jax_broadcast_in_dim_fn(arg1, arg2, broadcast_dims):
    return jax.lax.broadcast_in_dim(arg1, jax.numpy.shape(arg2), broadcast_dims)

broadcast_in_dim_symbolic_opinfo = OpInfo(
    lambda fd: partial(broadcast_in_dim_sym_fn, fd),
    "broadcast_in_dim_symbolic",
    sample_input_generator=broadcast_in_dim_generator,
    error_input_generator=broadcast_in_dim_error_generator,
    reference=jax_broadcast_in_dim_fn if JAX_AVAILABLE else None,
    reference_type=ReferenceType.Jax,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(broadcast_in_dim_symbolic_opinfo)


    error_input_generator=broadcast_in_dim_error_generator,
    reference=jax.lax.broadcast_in_dim if JAX_AVAILABLE else None,
    reference_type=ReferenceType.Jax,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        # This argument is purposely Constant even though the positional
        # argument can also be symbolic.
        ArgumentType.Constant,
        ArgumentType.Constant,
    ),
)
shape_ops.append(broadcast_in_dim_constant_opinfo)

# NOTE: The symbolic version of broadcast_in_dim opinfo tests the "shape"
# argument with a Vector generated from another operation like ops.shape.
def broadcast_in_dim_sym_fn(fd, arg1, arg2, broadcast_dims):
    return fd.ops.broadcast_in_dim(arg1, arg2.shape(), broadcast_dims)

def jax_broadcast_in_dim_fn(arg1, arg2, broadcast_dims):
    return jax.lax.broadcast_in_dim(arg1, jax.numpy.shape(arg2), broadcast_dims)

broadcast_in_dim_symbolic_opinfo = OpInfo(
    lambda fd: partial(broadcast_in_dim_sym_fn, fd),
    "broadcast_in_dim_symbolic",
    sample_input_generator=broadcast_in_dim_generator,
    error_input_generator=broadcast_in_dim_error_generator,
    reference=jax_broadcast_in_dim_fn if JAX_AVAILABLE else None,
    reference_type=ReferenceType.Jax,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(broadcast_in_dim_symbolic_opinfo)

# translate between nvfuser and pytorch argument order for scatter
def scatter_wrapper(


    return jax.lax.broadcast_in_dim(arg1, jax.numpy.shape(arg2), broadcast_dims)

broadcast_in_dim_symbolic_opinfo = OpInfo(
    lambda fd: partial(broadcast_in_dim_sym_fn, fd),
    "broadcast_in_dim_symbolic",
    sample_input_generator=broadcast_in_dim_generator,
    error_input_generator=broadcast_in_dim_error_generator,
    reference=jax_broadcast_in_dim_fn if JAX_AVAILABLE else None,
    reference_type=ReferenceType.Jax,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(broadcast_in_dim_symbolic_opinfo)

# translate between nvfuser and pytorch argument order for scatter
def scatter_wrapper(
    fn: callable, input: torch.Tensor, index: torch.Tensor, src: torch.Tensor, dim: int
):
    return fn(input, dim, index, src)

scatter_opinfo = OpInfo(
    lambda fd: fd.ops.scatter,
    "scatter",
    sample_input_generator=scatter_generator,
    reference=partial(scatter_wrapper, torch.scatter),
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(scatter_opinfo)

# translate between nvfuser and pytorch argument order for gather, take_along_dim, and index_select
def gather_wrapper(fn: callable, input: torch.Tensor, index: torch.Tensor, dim: int):
    return fn(input, dim, index)


):
    return fn(input, dim, index, src)

scatter_opinfo = OpInfo(
    lambda fd: fd.ops.scatter,
    "scatter",
    sample_input_generator=scatter_generator,
    reference=partial(scatter_wrapper, torch.scatter),
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(scatter_opinfo)

# translate between nvfuser and pytorch argument order for gather, take_along_dim, and index_select
def gather_wrapper(fn: callable, input: torch.Tensor, index: torch.Tensor, dim: int):
    return fn(input, dim, index)

gather_opinfo = OpInfo(
    lambda fd: fd.ops.gather,
    "gather",
    sample_input_generator=gather_generator,
    error_input_generator=take_along_axis_error_generator,
    reference=partial(gather_wrapper, torch.gather),
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(gather_opinfo)

index_select_opinfo = OpInfo(
    lambda fd: fd.ops.index_select,
    "index_select",
    sample_input_generator=index_select_generator,


    ),
)
shape_ops.append(gather_opinfo)

index_select_opinfo = OpInfo(
    lambda fd: fd.ops.index_select,
    "index_select",
    sample_input_generator=index_select_generator,
    error_input_generator=index_select_error_generator,
    reference=partial(gather_wrapper, torch.index_select),
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(index_select_opinfo)

# we needed a reference because argsort requires kwargs.
def argsort_ref(a, dim, descending, stable):
    return torch.argsort(a, dim=dim, descending=descending, stable=stable)

argsort_opinfo = OpInfo(
    lambda fd: fd.ops.argsort,
    "argsort",
    # TODO: complex dtypes are not supported by aten fallback
    dtypes=(int_float_dtypes),
    sample_input_generator=argsort_generator,
    reference=argsort_ref,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Constant,
        ArgumentType.Constant,
        ArgumentType.Constant,
    ),
)
shape_ops.append(argsort_opinfo)

topk_opinfo = OpInfo(


topk_opinfo = OpInfo(
    lambda fd: fd.ops.topk,
    "topk",
    dtypes=(int_float_dtypes),
    sample_input_generator=topk_generator,
    error_input_generator=topk_error_generator,
    reference=torch.topk,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,  # input tensor
        ArgumentType.Symbolic,  # k (number of elements)
        ArgumentType.Constant,  # dim
        ArgumentType.Constant,  # largest
        ArgumentType.Constant,  # sorted
    ),
)
shape_ops.append(topk_opinfo)

def index_put_accumulate_ref(
    acc: torch.Tensor, index: torch.Tensor, value: torch.Tensor
):
    return torch.index_put(
        acc,
        [
            index,
        ],
        value,
        accumulate=True,
    )

index_put_accumulate_opinfo = OpInfo(
    lambda fd: fd.ops.index_put_accumulate,
    "index_put_accumulate",
    sample_input_generator=index_put_accumulate_generator,
    reference=index_put_accumulate_ref,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
    ),
    # index_put_accumulate is not used in Thunder, so skip in direct bindings for now.
    supports_direct_bindings=False,
)
shape_ops.append(index_put_accumulate_opinfo)

# NvFuser's API is significantly different than JAX.
# TODO: Change python frontend api to match JAX using a cpp wrapper function.
pad_opinfo = OpInfo(


    reference=torch.permute,
    symbolic_parameter_list=(ArgumentType.Symbolic, ArgumentType.Constant),
)
shape_ops.append(permute_opinfo)

reshape_constant_opinfo = OpInfo(
    lambda fd: fd.ops.reshape,
    "reshape_constant",
    sample_input_generator=reshape_generator,
    error_input_generator=reshape_error_generator,
    reference=torch.reshape,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(reshape_constant_opinfo)

def reshape_sym_fn(fd, input_tensor, output_shaped_tensor):
    return fd.ops.reshape(input_tensor, output_shaped_tensor.shape())

def torch_reshape_sym_fn(input_tensor, output_shaped_tensor):
    return torch.reshape(input_tensor, output_shaped_tensor.size())

reshape_symbolic_opinfo = OpInfo(
    lambda fd: partial(reshape_sym_fn, fd),
    "reshape_symbolic",
    sample_input_generator=reshape_generator,
    error_input_generator=reshape_error_generator,
    reference=torch_reshape_sym_fn,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
    ),
)
shape_ops.append(reshape_symbolic_opinfo)


reshape_constant_opinfo = OpInfo(
    lambda fd: fd.ops.reshape,
    "reshape_constant",
    sample_input_generator=reshape_generator,
    error_input_generator=reshape_error_generator,
    reference=torch.reshape,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Constant,
    ),
)
shape_ops.append(reshape_constant_opinfo)

def reshape_sym_fn(fd, input_tensor, output_shaped_tensor):
    return fd.ops.reshape(input_tensor, output_shaped_tensor.shape())

def torch_reshape_sym_fn(input_tensor, output_shaped_tensor):
    return torch.reshape(input_tensor, output_shaped_tensor.size())

reshape_symbolic_opinfo = OpInfo(
    lambda fd: partial(reshape_sym_fn, fd),
    "reshape_symbolic",
    sample_input_generator=reshape_generator,
    error_input_generator=reshape_error_generator,
    reference=torch_reshape_sym_fn,
    symbolic_parameter_list=(
        ArgumentType.Symbolic,
        ArgumentType.Symbolic,
    ),
)
shape_ops.append(reshape_symbolic_opinfo)

slice_opinfo = OpInfo(
    lambda fd: fd.ops.slice,
    "slice",
    sample_input_generator=slice_generator,


        (torch.float16, torch.bfloat16)
        if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8
        else (torch.float16,)
    ),
    sample_input_generator=matmul_input_generator,
    reference=torch.matmul,
)
matmul_ops.append(matmul_opinfo)

# torch._grouped_mm and torch._scaled_grouped_mm is not available prior to PyTorch 2.8.0
if LooseVersion(torch.__version__) >= LooseVersion("2.8.0"):
    grouped_mm_opinfo = OpInfo(
        lambda fd: fd.ops.grouped_mm,
        "grouped_mm",
        # only bf16 is supported
        dtypes=(torch.bfloat16,),
        sample_input_generator=grouped_mm_input_generator,
        reference=torch._grouped_mm,
    )

    def scaled_grouped_mm_wrapper(
        mat1, mat2, offsets, scale1, scale2, alpha, bias, beta, dtype
    ):
        assert beta is None
        # mat1 needs to be in column major while mat2 needs to be in row major.
        row_major_mat2 = mat2.transpose(-1, -2).contiguous().transpose(-1, -2)
        if mat1.ndim == 2 and mat2.ndim == 2:
            # case 1, mat1 and mat2 are both 2D, aten fallback expects collapsed 1D scale with group dimension on the slower side.
            reshaped_scale1 = scale1.reshape(-1)
            reshaped_scale2 = scale2.reshape(-1)
        else:
            # squeeze out the k dimension
            reshaped_scale1 = scale1.squeeze(-1)
            reshaped_scale2 = scale2.squeeze(-2)
        return torch._scaled_grouped_mm(
            mat1,
            row_major_mat2,
            reshaped_scale1,
            reshaped_scale2,
            offsets,
            bias,
            alpha,
            dtype,
        )

    scaled_grouped_mm_opinfo = OpInfo(
        lambda fd: fd.ops.grouped_mm,
        "scaled_grouped_mm",
        # only float8 is supported
        dtypes=(torch.float8_e4m3fn,),
        sample_input_generator=scaled_grouped_mm_input_generator,
        reference=scaled_grouped_mm_wrapper,
        symbolic_parameter_list=(
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
        ),
    )


    scaled_grouped_mm_opinfo = OpInfo(
        lambda fd: fd.ops.grouped_mm,
        "scaled_grouped_mm",
        # only float8 is supported
        dtypes=(torch.float8_e4m3fn,),
        sample_input_generator=scaled_grouped_mm_input_generator,
        reference=scaled_grouped_mm_wrapper,
        symbolic_parameter_list=(
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
        ),
    )

    def scaled_mm_wrapper(mat1, mat2, scale1, scale2, alpha, bias, beta, dtype):
        assert beta is None
        return torch._scaled_mm(
            mat1, mat2, scale1, scale2, bias, alpha, out_dtype=dtype
        )

    scaled_mm_opinfo = OpInfo(
        lambda fd: fd.ops.scaled_mm,
        "scaled_mm",
        # limit test to mxfp8 for now
        dtypes=(torch.float8_e4m3fn,),
        sample_input_generator=scaled_mm_input_generator,
        reference=scaled_mm_wrapper,
        symbolic_parameter_list=(
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Symbolic,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
            ArgumentType.Constant,
        ),
    )