Jhnony/Mobilekernelbench huggingface dataset:
# Pooling / GlobalMaxPool
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class Model(nn.Module):
    """
    A model that performs global max pooling operation.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global max pooling across all spatial dimensions.
        Each channel is reduced to a single value by taking the maximum.
        
        For 2D input (N, C, H, W) -> output (N, C, 1, 1)
        For 3D input (N, C, D, H, W) -> output (N, C, 1, 1, 1)
        
        Y[n, c] = max(X[n, c, :, :, ...]) over all spatial dimensions
    """
    
    def __init__(self):
        """
        Initialize the GlobalMaxPool model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor of shape (N, C, H, W) for 2D or (N, C, D, H, W) for 3D
        
        Returns:
            Output tensor of shape (N, C, 1, 1) for 2D or (N, C, 1, 1, 1) for 3D
        
        Raises:
            ValueError: If input dimension is not 4D or 5D
        """
        if x.dim() == 4:
            # 2D Global Max Pooling: (N, C, H, W) -> (N, C, 1, 1)
            return F.adaptive_max_pool2d(x, output_size=1)
        elif x.dim() == 5:
            # 3D Global Max Pooling: (N, C, D, H, W) -> (N, C, 1, 1, 1)
            return F.adaptive_max_pool3d(x, output_size=1)
        else:
            raise ValueError(
                f"GlobalMaxPool expects 4D or 5D input, "
                f"but got {x.dim()}D input with shape {x.shape}"
            )


# ======== Example input configuration ========

batch_size = 16
channels = 3
height = 128
width = 128

def get_inputs():
    torch.manual_seed(42)
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return []
# Pooling / GlobalMaxPool_2d_single_channel
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class Model(nn.Module):
    """
    A model that performs global max pooling operation.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global max pooling across all spatial dimensions.
        Each channel is reduced to a single value by taking the maximum.
        
        For 2D input (N, C, H, W) -> output (N, C, 1, 1)
        For 3D input (N, C, D, H, W) -> output (N, C, 1, 1, 1)
        
        Y[n, c] = max(X[n, c, :, :, ...]) over all spatial dimensions
    """
    
    def __init__(self):
        """
        Initialize the GlobalMaxPool model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor of shape (N, C, H, W) for 2D or (N, C, D, H, W) for 3D
        
        Returns:
            Output tensor of shape (N, C, 1, 1) for 2D or (N, C, 1, 1, 1) for 3D
        
        Raises:
            ValueError: If input dimension is not 4D or 5D
        """
        if x.dim() == 4:
            # 2D Global Max Pooling: (N, C, H, W) -> (N, C, 1, 1)
            return F.adaptive_max_pool2d(x, output_size=1)
        elif x.dim() == 5:
            # 3D Global Max Pooling: (N, C, D, H, W) -> (N, C, 1, 1, 1)
            return F.adaptive_max_pool3d(x, output_size=1)
        else:
            raise ValueError(
                f"GlobalMaxPool expects 4D or 5D input, "
                f"but got {x.dim()}D input with shape {x.shape}"
            )


# ======== Example input configuration ========

batch_size = 32
channels = 3
height = 128
width = 64

def get_inputs():
    torch.manual_seed(42)
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return []
# Pooling / GlobalMaxPool_3d
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class Model(nn.Module):
    """
    A model that performs 3D global max pooling for video or medical imaging.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global max pooling across all spatial and temporal dimensions.
        Used for video classification, action recognition, or 3D medical imaging.
        
        Input: (N, C, D, H, W) -> Output: (N, C, 1, 1, 1)
        Y[n, c, 1, 1, 1] = max(X[n, c, d, h, w]) for all d, h, w
    """
    
    def __init__(self):
        """
        Initialize the 3D GlobalMaxPool model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor of shape (N, C, D, H, W)
        
        Returns:
            Output tensor of shape (N, C, 1, 1, 1)
        """
        return F.adaptive_max_pool3d(x, output_size=1)


# ======== Example input configuration ========

batch_size = 8
channels = 256
depth = 16
height = 8
width = 8

def get_inputs():
    torch.manual_seed(42)
    x = torch.randn(batch_size, channels, depth, height, width)
    return [x]

def get_init_inputs():
    return []
# Pooling / MaxPool
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        super(Model, self).__init__()
        self.max_pool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.max_pool(x)

batch_size = 2
channels = 16
height = 64
width = 64
kernel_size = 11

def get_inputs():
    x = torch.rand(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return [kernel_size]
# Pooling / MaxPool_2d_ceil
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional


class Model(nn.Module):
    """
    A model that applies 2D Max Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies max pooling operation with specified kernel size, stride, padding, and ceil mode.
    """
    
    def __init__(self, kernel_shape: Tuple[int, int] = (3, 3), 
                 strides: Tuple[int, int] = (2, 2), 
                 pads: Optional[Tuple[int, int, int, int]] = None,
                 ceil_mode: int = 1):
        super(Model, self).__init__()
        self.kernel_shape = kernel_shape
        self.strides = strides
        self.pads = pads
        self.ceil_mode = bool(ceil_mode)  # Convert to boolean for PyTorch
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, channels, height, width]
        
        Returns:
            Output tensor after max pooling
        """
        # Convert pads to PyTorch padding format if provided
        # ONNX uses [top, left, bottom, right], PyTorch uses (left, right, top, bottom)
        if self.pads is not None:
            x = F.pad(x, (self.pads[1], self.pads[3], self.pads[0], self.pads[2]))
            padding = 0
        else:
            padding = 0
        
        return F.max_pool2d(x, kernel_size=self.kernel_shape, 
                           stride=self.strides, 
                           padding=padding,
                           ceil_mode=self.ceil_mode)


# ======== Example input configuration ========

batch_size = 1
channels = 3
height = 32
width = 32

def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return [(3, 3), (2, 2), None, 1]  # [kernel_shape, strides, pads, ceil_mode]
# Pooling / MaxPool_2d_dilations
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional


class Model(nn.Module):
    """
    A model that applies 2D Max Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies max pooling operation with specified kernel size, stride, padding, and dilation.
    """
    
    def __init__(self, kernel_shape: Tuple[int, int] = (2, 2), 
                 strides: Tuple[int, int] = (1, 1), 
                 pads: Optional[Tuple[int, int, int, int]] = None,
                 dilations: Tuple[int, int] = (2, 2)):
        super(Model, self).__init__()
        self.kernel_shape = kernel_shape
        self.strides = strides
        self.pads = pads
        self.dilations = dilations
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, channels, height, width]
        
        Returns:
            Output tensor after max pooling
        """
        # Convert pads to PyTorch padding format if provided
        # ONNX uses [top, left, bottom, right], PyTorch uses (left, right, top, bottom)
        if self.pads is not None:
            x = F.pad(x, (self.pads[1], self.pads[3], self.pads[0], self.pads[2]))
            padding = 0
        else:
            padding = 0
        
        return F.max_pool2d(x, kernel_size=self.kernel_shape, 
                           stride=self.strides, 
                           padding=padding,
                           dilation=self.dilations)


# ======== Example input configuration ========

batch_size = 1
channels = 3
height =32
width = 32

def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return [(2, 2), (1, 1), None, (2, 2)]  # [kernel_shape, strides, pads, dilations]
# Pooling / MaxPool_2d_same_lower
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional


class Model(nn.Module):
    """
    A model that applies 2D Max Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies max pooling operation with specified kernel size, stride, and padding.
    """
    
    def __init__(self, kernel_shape: Tuple[int, int] = (2, 2), 
                 strides: Tuple[int, int] = (1, 1), 
                 pads: Optional[Tuple[int, int, int, int]] = None,
                 auto_pad: str = "SAME_LOWER"):
        super(Model, self).__init__()
        self.kernel_shape = kernel_shape
        self.strides = strides
        self.pads = pads
        self.auto_pad = auto_pad
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, channels, height, width]
        
        Returns:
            Output tensor after max pooling
        """
        # Handle auto_pad
        if self.auto_pad in ["SAME_UPPER", "SAME_LOWER"]:
            # Calculate padding for 'SAME' mode
            padding = self._calculate_same_padding(x.shape[2:])
        elif self.pads is not None:
            # ONNX uses [top, left, bottom, right], PyTorch uses (left, right, top, bottom)
            x = F.pad(x, (self.pads[1], self.pads[3], self.pads[0], self.pads[2]))
            padding = 0
        else:
            padding = 0
        
        return F.max_pool2d(x, kernel_size=self.kernel_shape, 
                           stride=self.strides, padding=padding)
    
    def _calculate_same_padding(self, input_shape: Tuple[int, int]) -> Tuple[int, int]:
        """Calculate padding for SAME mode"""
        padding = []
        for i in range(2):
            out_size = (input_shape[i] + self.strides[i] - 1) // self.strides[i]
            pad_total = max(0, (out_size - 1) * self.strides[i] + self.kernel_shape[i] - input_shape[i])
            padding.append(pad_total // 2)
        return tuple(padding)


# ======== Example input configuration ========

batch_size = 1
channels = 3
height = 32
width = 32

def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return [(2, 2), (1, 1), None, "SAME_LOWER"]  # [kernel_shape, strides, pads, auto_pad]
# Pooling / MaxPool_2d_strides
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Union, Optional


class Model(nn.Module):
    """
    A model that applies 2D Max Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies max pooling operation with specified kernel size, stride, and padding.
    """
    
    def __init__(self, kernel_shape: Tuple[int, int] = (5, 5), 
                 strides: Tuple[int, int] = (3, 3), 
                 pads: Optional[Tuple[int, int, int, int]] = None):
        super(Model, self).__init__()
        self.kernel_shape = kernel_shape
        self.strides = strides
        self.pads = pads
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, channels, height, width]
        
        Returns:
            Output tensor after max pooling
        """
        # Convert pads to PyTorch padding format if provided
        # ONNX uses [top, left, bottom, right], PyTorch uses (left, right, top, bottom)
        if self.pads is not None:
            # Apply padding manually if asymmetric padding is needed
            x = F.pad(x, (self.pads[1], self.pads[3], self.pads[0], self.pads[2]))
            padding = 0
        else:
            padding = 0
        
        return F.max_pool2d(x, kernel_size=self.kernel_shape, 
                           stride=self.strides, padding=padding)


# ======== Example input configuration ========

batch_size = 1
channels = 3
height = 32
width = 32

def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]

def get_init_inputs():
    return [(5, 5), (3, 3), None]  # [kernel_shape, strides, pads]
# Reduction / ArgMax
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Argmax over a specified dimension.
    """
    def __init__(self, dim: int):
        """
        Initializes the model with the dimension to perform argmax.

        Args:
            dim (int): The dimension to perform argmax over.
        """
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies argmax over the specified dimension to the input tensor.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor with argmax applied, with the specified dimension removed.
        """
        return torch.argmax(x, dim=self.dim)

batch_size = 32
dim1 = 128
dim2 = 256

def get_inputs():
    x = torch.rand(batch_size, dim1, dim2)
    return [x]

def get_init_inputs():
    return [1]
# Reduction / ArgMax_default_axis_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmax(data)

batch_size = 16
dim1 = 128
dim2 = 256

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]

def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ArgMax_keepdims_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.dim = 0

    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmax(data, dim=self.dim, keepdim=True)

batch_size = 32
dim1 = 128
dim2 = 512

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]

def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ArgMax_negative_axis_keepdims_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.dim = -1
    
    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmax(data, dim=self.dim, keepdim=True)

batch_size = 16
dim1 = 128
dim2 = 256

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ArgMin
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that finds the index of the minimum value along a specified dimension.
    """
    def __init__(self, dim: int):
        """
        Initializes the model with the dimension to perform argmin on.

        Args:
            dim (int): Dimension along which to find the minimum value.
        """
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Finds the index of the minimum value along the specified dimension.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Tensor containing the indices of the minimum values along the specified dimension.
        """
        return torch.argmin(x, dim=self.dim)

batch_size = 32
dim1 = 128
dim2 = 256
dim = 1

def get_inputs():
    x = torch.rand(batch_size, dim1, dim2)
    return [x]

def get_init_inputs():
    return [dim]
# Reduction / ArgMin_default_axis_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmin(data)

batch_size = 16
dim1 = 128
dim2 = 127

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]

def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ArgMin_keepdims_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.dim = 0
    
    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmin(data, dim=self.dim, keepdim=True)

batch_size = 16
dim1 = 256
dim2 = 256

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ArgMin_negative_axis_keepdims_example
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.dim = -1
    
    def forward(self, data) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            data: Tensor
        
        Returns:
            Output tensor
        """
        return torch.argmin(data, dim=self.dim, keepdim=True)

batch_size = 16
dim1 = 128
dim2 = 256

def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim1, dim2),  # data
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return []
# Reduction / ReduceL1_2d_input
import torch
import torch.nn as nn
from typing import Optional

class Model(nn.Module):
    """
    ReduceL1 on 2D input
    
    Common for matrix operations.
    
    Example:
        >>> model = Model(axes=1, keepdims=0)
        >>> x = torch.tensor([[-1.0, 2.0, -3.0], [4.0, -5.0, 6.0]])
        >>> axes_tensor = torch.tensor([1])
        >>> output = model(x, axes_tensor)
        >>> print(output)  # tensor([6., 15.])  # Row-wise L1 norms
    """
    
    def __init__(self, axes: int, keepdims: int = 0):
        super(Model, self).__init__()
        self.axes = axes
        self.keepdims = bool(keepdims)
    
    def forward(self, data: torch.Tensor, axes: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Compute L1 norm over 2D tensor
        
        Args:
            data: 2D input tensor
            axes: Axes to reduce (optional)
            
        Returns:
            Tensor with specified dimension reduced
        """
        if axes is None:
            axes = self.axes
        else:
            axes = axes.item() if axes.numel() == 1 else axes.tolist()
        
        return torch.norm(data, p=1, dim=axes, keepdim=self.keepdims)

rows = 64
cols = 128
axes = 1
keepdims = 0

def get_inputs():
    torch.manual_seed(42)
    data = torch.randn(rows, cols)
    return [data]

def get_init_inputs():
    return [axes, keepdims]
# Reduction / ReduceL1_multiple_axes
import torch
import torch.nn as nn
from typing import Tuple, Optional

class Model(nn.Module):
    """
    ReduceL1 over multiple axes
    
    Reduces along multiple dimensions simultaneously.
    
    Example:
        >>> model = Model(axes=(0, 2), keepdims=0)
        >>> x = torch.randn(2, 3, 4)
        >>> axes_tensor = torch.tensor([0, 2])
        >>> output = model(x, axes_tensor)
        >>> print(output.shape)  # torch.Size([3])
    """
    
    def __init__(self, axes: Tuple[int, ...], keepdims: int = 0):
        super(Model, self).__init__()
        self.axes = axes
        self.keepdims = bool(keepdims)
    
    def forward(self, data: torch.Tensor, axes: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Compute L1 norm over multiple specified dimensions
        
        Args:
            data: Input tensor
            axes: Axes to reduce (optional)
            
        Returns:
            Tensor with L1 norm over multiple dimensions
        """
        if axes is None:
            axes_list = self.axes
        else:
            axes_list = tuple(axes.tolist())
        
        return torch.norm(data, p=1, dim=axes_list, keepdim=self.keepdims)

axes = (0, 2)
keepdims = 0

def get_inputs():
    torch.manual_seed(42)
    data = torch.randn(2, 3, 4)
    return [data]

def get_init_inputs():
    return [axes, keepdims]