Jhnony/Mobilekernelbench huggingface dataset:
# Normalization / LayerNorm
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Layer Normalization.
    """
    def __init__(self, normalized_shape: tuple):
        """
        Initializes the LayerNorm layer.

        Args:
            normalized_shape (tuple): Shape of the input tensor to be normalized.
        """
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Layer Normalization to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (*, normalized_shape).

        Returns:
            torch.Tensor: Output tensor with Layer Normalization applied, same shape as input.
        """
        return self.ln(x)

batch_size = 4
features = 16
dim1 = 64
dim2 = 64

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

def get_init_inputs():
    return [(features, dim1, dim2)]
# Normalization / LayerNorm_1d
import torch
import torch.nn as nn
from typing import List, Tuple


class Model(nn.Module):
    """
    A model that applies Layer Normalization over a single dimension.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes input over the last dimension.
        Commonly used in transformers and sequence models.
    """
    
    def __init__(self, normalized_shape: int, eps: float = 1e-5, elementwise_affine: bool = True):
        """
        Args:
            normalized_shape: Size of the dimension to normalize
            eps: Small value added to variance for numerical stability
            elementwise_affine: If True, learnable affine parameters
        """
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(
            normalized_shape=normalized_shape,
            eps=eps,
            elementwise_affine=elementwise_affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, seq_len, features]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.ln(x)


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

batch_size = 32
seq_len = 128
features = 512
eps = 1e-5
elementwise_affine = True

def get_inputs():
    x = torch.randn(batch_size, seq_len, features)
    return [x]

def get_init_inputs():
    return [features, eps, elementwise_affine]
# Normalization / LayerNorm_3d
import torch
import torch.nn as nn
from typing import List, Tuple


class Model(nn.Module):
    """
    A model that applies Layer Normalization over last three dimensions.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes input over the last three dimensions.
        Useful for 4D tensors (batch, channels, height, width).
    """
    
    def __init__(self, normalized_shape: Tuple[int, int, int], eps: float = 1e-5,
                 elementwise_affine: bool = True):
        """
        Args:
            normalized_shape: Tuple of (channels, height, width) to normalize
            eps: Small value added to variance for numerical stability
            elementwise_affine: If True, learnable affine parameters
        """
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(
            normalized_shape=normalized_shape,
            eps=eps,
            elementwise_affine=elementwise_affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, channels, height, width]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.ln(x)


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

batch_size = 4
channels = 16
height = 32
width = 32
eps = 1e-5
elementwise_affine = True

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

def get_init_inputs():
    return [(channels, height, width), eps, elementwise_affine]
# Normalization / LayerNorm_eps
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that applies Layer Normalization with higher epsilon value.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes input with a larger epsilon for stability.
        Tests numerical stability with different epsilon values.
    """
    
    def __init__(self, normalized_shape: int, eps: float = 1e-3, 
                 elementwise_affine: bool = True):
        """
        Args:
            normalized_shape: Size of the dimension to normalize
            eps: Small value added to variance for numerical stability
            elementwise_affine: If True, learnable affine parameters
        """
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(
            normalized_shape=normalized_shape,
            eps=eps,
            elementwise_affine=elementwise_affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, seq_len, features]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.ln(x)


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

batch_size = 32
seq_len = 128
features = 256
eps = 1e-3
elementwise_affine = True

def get_inputs():
    x = torch.randn(batch_size, seq_len, features)
    return [x]

def get_init_inputs():
    return [features, eps, elementwise_affine]
# Normalization / LayerNorm_no_affine
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class Model(nn.Module):
    """
    A model that applies Layer Normalization without learnable parameters.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies Layer Normalization without affine transformation.
        Only normalizes to zero mean and unit variance.
    """
    
    def __init__(self, normalized_shape: int, eps: float = 1e-5):
        super(Model, self).__init__()
        self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps, elementwise_affine=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after layer normalization (no affine)
        """
        return self.layer_norm(x)


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

batch_size = 8
seq_len = 256
hidden_dim = 512

def get_inputs():
    x = torch.randn(batch_size, seq_len, hidden_dim)
    return [x]

def get_init_inputs():
    return [hidden_dim]
# Normalization / LayerNorm_vision
import torch
import torch.nn as nn
from typing import List, Tuple


class Model(nn.Module):
    """
    A model that applies Layer Normalization for vision transformers.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes patches in vision transformer style.
        Normalizes over (num_patches, embed_dim).
    """
    
    def __init__(self, normalized_shape: Tuple[int, int], eps: float = 1e-6, 
                 elementwise_affine: bool = True):
        """
        Args:
            normalized_shape: Tuple of (num_patches, embed_dim) to normalize
            eps: Small value added to variance for numerical stability
            elementwise_affine: If True, learnable affine parameters
        """
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(
            normalized_shape=normalized_shape,
            eps=eps,
            elementwise_affine=elementwise_affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, num_patches, embed_dim]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.ln(x)


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

batch_size = 16
num_patches = 196  # 14x14 patches for 224x224 image with 16x16 patch size
embed_dim = 768
eps = 1e-6
elementwise_affine = True

def get_inputs():
    x = torch.randn(batch_size, num_patches, embed_dim)
    return [x]

def get_init_inputs():
    return [(num_patches, embed_dim), eps, elementwise_affine]
# Others / LSTM
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that uses LSTM (Long Short-Term Memory) layer.

    Semantics:
        output, (h_n, c_n) = LSTM(x, (h_0, c_0))
        Applies LSTM to an input sequence.
        Input shape: (seq_length, batch_size, input_size)
        Output shape: (seq_length, batch_size, hidden_size)
    """

    def __init__(self, input_size=16, hidden_size=32):
        super(Model, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=False,
            bidirectional=False
        )

    def forward(self, x: torch.Tensor, h_0: torch.Tensor, c_0: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input tensor of shape (seq_length, batch_size, input_size)
            h_0: Initial hidden state of shape (1, batch_size, hidden_size)
            c_0: Initial cell state of shape (1, batch_size, hidden_size)
        
        Returns:
            output: Tensor of shape (seq_length, batch_size, hidden_size)
        """
        output, (h_n, c_n) = self.lstm(x, (h_0, c_0))
        return output


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

batch_size = 4
seq_length = 8
input_size = 16
hidden_size = 32

def get_inputs():
    """
    Generate random input sequence and initial states for LSTM.
    """
    x = torch.randn(seq_length, batch_size, input_size)
    h_0 = torch.randn(1, batch_size, hidden_size)
    c_0 = torch.randn(1, batch_size, hidden_size)
    return [x, h_0, c_0]

def get_init_inputs():
    return []
# Others / RNN
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that uses RNN (Recurrent Neural Network) layer.

    Semantics:
        output, h_n = RNN(x, h_0)
        Applies a multi-layer Elman RNN to an input sequence.
        Input shape: (seq_length, batch_size, input_size)
        Output shape: (seq_length, batch_size, hidden_size)
    """

    def __init__(self):
        super(Model, self).__init__()
        self.rnn = nn.RNN(input_size=10, hidden_size=20, num_layers=1, batch_first=False)

    def forward(self, x: torch.Tensor, h_0: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input tensor of shape (seq_length, batch_size, input_size)
            h_0: Initial hidden state of shape (1, batch_size, hidden_size)
        
        Returns:
            output: Tensor of shape (seq_length, batch_size, hidden_size)
        """
        output, h_n = self.rnn(x, h_0)
        return output


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

batch_size = 4
seq_length = 8
input_size = 10
hidden_size = 20

def get_inputs():
    x = torch.randn(seq_length, batch_size, input_size)
    h_0 = torch.randn(1, batch_size, hidden_size)
    return [x, h_0]

def get_init_inputs():
    return []
# Others / ROIAlign
import torch
import torch.nn as nn
import torchvision

class Model(nn.Module):
    """
    Simple model that performs ROI Align.
    """
    def __init__(self, output_size: tuple, spatial_scale: float, sampling_ratio: int, aligned: bool):
        super(Model, self).__init__()
        self.output_size = output_size
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        self.aligned = aligned

    def forward(self, x: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
        """
        Applies RoiAlign.
        """
        return torchvision.ops.roi_align(
            x, rois, 
            output_size=self.output_size, 
            spatial_scale=self.spatial_scale, 
            sampling_ratio=self.sampling_ratio,
            aligned=self.aligned
        )

batch_size = 8
channels = 16
height = 128
width = 128

def get_inputs():
    x = torch.rand(batch_size, channels, height, width)
    # ROIs: [batch_idx, x1, y1, x2, y2]
    rois = torch.tensor([
        [0, 2.5, 2.5, 12.5, 12.5],
        [1, 8.0, 8.0, 24.0, 24.0]
    ])
    return [x, rois]

def get_init_inputs():
    # output_size=(7, 7), spatial_scale=1.0, sampling_ratio=2, aligned=True
    return [(7, 7), 1.0, 2, True]
# Others / Stft
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, n_fft, hop_length, win_length, window=None, onesided=True, center=True):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.onesided = onesided
        self.center = center
        
        if window is not None:
            self.register_buffer("window", window)
        else:
            self.register_buffer("window", torch.hann_window(win_length))

    def forward(self, x):
        """
        x: [N, L] or [L]
        Returns: STFT result in MNN format: [N, F, T, 2]
        (real+imag split)
        """

        # PyTorch STFT
        stft = torch.stft(
            x,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            center=self.center,
            return_complex=False,
            onesided=self.onesided
        )

        # PyTorch outputs: [*, freq, frames, 2]
        return stft

n_fft = 256
hop_length = 64
win_length = 200

# ======== Example input configuration ========
def get_inputs() -> list:
    x = torch.randn(4, 16000) * 0.1
    return [x]

def get_init_inputs():
    return [n_fft, hop_length, win_length]
# Pooling / AveragePool
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs 2D Average Pooling.
    """
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        """
        Initializes the Average Pooling layer.

        Args:
            kernel_size (int): Size of the pooling window.
            stride (int, optional): Stride of the pooling operation. Defaults to None (same as kernel_size).
            padding (int, optional): Padding applied to the input tensor. Defaults to 0.
        """
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies 2D Average Pooling to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied.
        """
        return self.avg_pool(x)

batch_size = 32
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 / AveragePool_2d_ceil
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):
    """
    Model that performs 2D Average Pooling with ceil mode.
    """
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=True)
    
    def forward(self, x) -> torch.Tensor:
        """
        Applies 2D Average Pooling with ceil mode to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied.
        """
        return self.avg_pool(x)

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

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, channels, height, width),  # x
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return [kernel_size]
# Pooling / AveragePool_2d_strides
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):
    """
    Model that performs 2D Average Pooling with custom strides.
    """
    def __init__(self, kernel_size: int, stride: int, padding: int = 0):
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=padding)
    
    def forward(self, x) -> torch.Tensor:
        """
        Applies 2D Average Pooling with custom strides to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied.
        """
        return F.avg_pool2d(x, kernel_size=(5, 5), stride=(3, 3), padding=None[:2] if None else 0)

batch_size = 8
channels = 32
height = 64
width = 64
kernel_size = 11
stride = 3

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, channels, height, width),  # x
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return [kernel_size, stride]
# Pooling / GlobalAveragePool
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class Model(nn.Module):
    """
    A model that applies Global Average Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global average pooling across all spatial dimensions.
        Each channel is reduced to a single value by computing the mean
        of all spatial locations.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [N, C, H, W] for 2D or [N, C, D, H, W] for 3D
        
        Returns:
            Output tensor with shape [N, C, 1, 1] for 2D or [N, C, 1, 1, 1] for 3D
        """
        if x.dim() == 4:
            # 2D Global Average Pooling: (N, C, H, W) -> (N, C, 1, 1)
            return F.adaptive_avg_pool2d(x, output_size=1)
        elif x.dim() == 5:
            # 3D Global Average Pooling: (N, C, D, H, W) -> (N, C, 1, 1, 1)
            return F.adaptive_avg_pool3d(x, output_size=1)
        else:
            raise ValueError(f"Expected 4D or 5D input, got {x.dim()}D")


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

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

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

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


class Model(nn.Module):
    """
    A model that applies 2D Global Average Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global average pooling across spatial dimensions (H, W).
        Output shape: (N, C, 1, 1)
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [N, C, H, W]
        
        Returns:
            Output tensor with shape [N, C, 1, 1]
        """
        return F.adaptive_avg_pool2d(x, output_size=1)


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

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

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

def get_init_inputs():
    return []
# Pooling / GlobalAveragePool_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 applies 2D Global Average Pooling to a single-channel input.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global average pooling across spatial dimensions (H, W).
        Tests with single channel input.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [N, 1, H, W]
        
        Returns:
            Output tensor with shape [N, 1, 1, 1]
        """
        return F.adaptive_avg_pool2d(x, output_size=1)


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

batch_size = 16
channels = 3
height = 64
width = 64

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

def get_init_inputs():
    return []
# Pooling / GlobalAveragePool_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 applies 3D Global Average Pooling to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies global average pooling across spatial dimensions (D, H, W).
        Output shape: (N, C, 1, 1, 1)
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [N, C, D, H, W]
        
        Returns:
            Output tensor with shape [N, C, 1, 1, 1]
        """
        return F.adaptive_avg_pool3d(x, output_size=1)


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

batch_size = 16
channels = 3
depth = 16
height = 64
width = 64

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

def get_init_inputs():
    return []