Jhnony/Mobilekernelbench huggingface dataset:
# Convolution / ConvTranspose_dilations
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Tuple


class Model(nn.Module):
    """
    A model that applies 2D Transposed Convolution to the input tensor.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Applies transposed convolution operation with specified stride, padding, dilation, groups, and output_padding.
    """
    
    def __init__(self, strides: Tuple[int, int] = (1, 1),
                 pads: Optional[Tuple[int, int, int, int]] = None,
                 dilations: Tuple[int, int] = (2, 2),
                 group: int = 1,
                 output_padding: Tuple[int, int] = (0, 0)):
        super(Model, self).__init__()
        self.strides = strides
        self.pads = pads
        self.dilations = dilations
        self.group = group
        self.output_padding = output_padding
    
    def forward(self, X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            X: Input tensor with shape [batch_size, in_channels, height, width]
            W: Weight tensor with shape [in_channels, out_channels/groups, kernel_h, kernel_w]
        
        Returns:
            Output tensor after transposed convolution
        """
        # 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.conv_transpose2d(X, W, bias=None, stride=self.strides, 
                                 padding=padding, output_padding=self.output_padding,
                                 groups=self.group, dilation=self.dilations)


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

batch_size = 4
in_channels = 16
height = 32
width = 32
out_channels = 4
kernel_h = 32
kernel_w = 32

def get_inputs():
    X = torch.randn(batch_size, in_channels, height, width)
    W = torch.randn(in_channels, out_channels, kernel_h, kernel_w)
    return [X, W]

def get_init_inputs():
    return [(1, 1), None, (2, 2), 1, (0, 0)]  # [strides, pads, dilations, group, output_padding]
# Convolution / Conv_with_strides_no_padding
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple


class Model(nn.Module):
    """
    A model that applies 2D Convolution to the input tensor.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Applies 2D convolution operation with specified stride, padding, dilation, and groups.
    """
    
    def __init__(self, strides: Tuple[int, int] = (2, 2),
                 pads: Tuple[int, int, int, int] = (0, 0, 0, 0),
                 dilations: Tuple[int, int] = (1, 1),
                 group: int = 1,
                 kernel_shape: Tuple[int, int] = (3, 3)):
        super(Model, self).__init__()
        self.strides = strides
        self.pads = pads
        self.dilations = dilations
        self.group = group
        self.kernel_shape = kernel_shape
    
    def forward(self, x: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, in_channels, height, width]
            W: Weight tensor with shape [out_channels, in_channels/groups, kernel_h, kernel_w]
        
        Returns:
            Output tensor after convolution
        """
        # Convert ONNX padding format [top, left, bottom, right] to PyTorch format (left, right)
        # PyTorch conv2d uses symmetric padding (left/right, top/bottom)
        padding = (self.pads[0], self.pads[1])
        
        return F.conv2d(x, W, bias=None, stride=self.strides, 
                       padding=padding, dilation=self.dilations, groups=self.group)


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


batch_size = 8
in_channels = 8
height = 64
width = 64
out_channels = 4
kernel_h = 64
kernel_w = 64

def get_inputs():
    x = torch.randn(batch_size, in_channels, height, width)
    W = torch.randn(out_channels, in_channels, kernel_h, kernel_w)
    return [x, W]

def get_init_inputs():
    return [(2, 2), (0, 0, 0, 0), (1, 1), 1, (3, 3)]  # [strides, pads, dilations, group, kernel_shape]
# Convolution / Conv_with_strides_padding
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple


class Model(nn.Module):
    """
    A model that applies 2D Convolution to the input tensor.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Applies 2D convolution operation with specified stride, padding, dilation, and groups.
    """
    
    def __init__(self, strides: Tuple[int, int] = (2, 2),
                 pads: Tuple[int, int, int, int] = (1, 1, 1, 1),
                 dilations: Tuple[int, int] = (1, 1),
                 group: int = 1,
                 kernel_shape: Tuple[int, int] = (3, 3)):
        super(Model, self).__init__()
        self.strides = strides
        self.pads = pads
        self.dilations = dilations
        self.group = group
        self.kernel_shape = kernel_shape
    
    def forward(self, x: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, in_channels, height, width]
            W: Weight tensor with shape [out_channels, in_channels/groups, kernel_h, kernel_w]
        
        Returns:
            Output tensor after convolution
        """
        # Convert ONNX padding format [top, left, bottom, right] to PyTorch format (left, right)
        # PyTorch conv2d uses symmetric padding (left/right, top/bottom)
        padding = (self.pads[0], self.pads[1])
        
        return F.conv2d(x, W, bias=None, stride=self.strides, 
                       padding=padding, dilation=self.dilations, groups=self.group)


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

batch_size = 8
in_channels = 8
height = 64
width = 64
out_channels = 4
kernel_h = 64
kernel_w = 64

def get_inputs():
    x = torch.randn(batch_size, in_channels, height, width)
    W = torch.randn(out_channels, in_channels, kernel_h, kernel_w)
    return [x, W]

def get_init_inputs():
    return [(2, 2), (1, 1, 1, 1), (1, 1), 1, (3, 3)]  # [strides, pads, dilations, group, kernel_shape]
# Convolution / DeconvolutionDepthwise_2D_kernel
import torch
import torch.nn as nn


class DepthwiseDeconv2dK5(nn.Module):
    """
    Depthwise deconvolution (ConvTranspose2d) with kernel_size=5, stride=1, padding=2.
    """
    def __init__(self, in_ch: int , out_ch: int , k: int ,stride, padding, output_padding, dilation):
        super().__init__()
        self.deconv = nn.ConvTranspose2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=k,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
            dilation=dilation,
            groups=in_ch,  # depthwise
            bias=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.deconv(x)
# Parameters
batch_size = 8
in_channels = 32
out_channels = 32  # For depthwise, typically in_channels == out_channels == groups
height = 48
width = 48
kernel_size = 5
stride = 1
padding = 2
output_padding = 0
dilation = 1
def get_inputs():
    x = torch.randn(batch_size, in_channels, height, width)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, dilation]
# Convolution / DeconvolutionDepthwise_2D_stride
import torch
import torch.nn as nn


class DepthwiseDeconv2dStride2(nn.Module):
    """
    Depthwise deconvolution (ConvTranspose2d) with stride=2 and padding=1.
    """
    def __init__(self, in_ch: int, out_ch: int , k: int, stride, padding, output_padding):
        super().__init__()
        self.deconv = nn.ConvTranspose2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=k,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
            groups=in_ch,  # depthwise
            bias=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.deconv(x)
# Parameters
batch_size = 8
in_channels = 64
out_channels = 64  # For depthwise, typically in_channels == out_channels == groups
height = 64
width = 64
kernel_size = 3
stride = 2
padding = 1
output_padding = 1

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding]
# Convolution / Deconvolution_1D
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 1D convolution operation.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        output_padding (int, optional): Additional size added to one side of the output shape. Defaults to 0.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 1D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, length).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, length_out).
        """
        return self.conv1d_transpose(x)

# Test code
batch_size = 8
in_channels = 64
out_channels = 64
kernel_size = 3
# much larger signal length for heavier workload
length = 1024

def get_inputs():
    x = torch.rand(batch_size, in_channels, length)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization
# Convolution / Deconvolution_2D_asymmetric_input_square_kernel___dilated____padded____strided__
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 2D transposed convolution operation with asymmetric input and square kernel, supporting dilation, padding, and stride.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel (square, e.g., 3 for a 3x3 kernel).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 2D transposed convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height_in, width_in). 

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)


# Test code
batch_size = 2
in_channels = 16
out_channels = 32
kernel_size = 3
height_in = 32
width_in = 64
stride = 5
padding = 1
dilation = 2

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, dilation]
# Convolution / Deconvolution_2D__asymmetric_input__asymmetric_kernel
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 2D convolution operation with asymmetric input and kernel size.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Tuple of integers representing the kernel size (height, width).
        stride (tuple, optional): Tuple of integers representing the stride of the convolution. Defaults to (1, 1).
        padding (tuple, optional): Tuple of integers representing the padding applied to the input. Defaults to (0, 0).
        output_padding (tuple, optional): Tuple of integers representing the additional size added to one side of the output shape. Defaults to (0, 0).
        dilation (tuple, optional): Tuple of integers representing the spacing between kernel elements. Defaults to (1, 1).
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), output_padding: tuple = (0, 0), dilation: tuple = (1, 1), groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height_in, width_in).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)

# Test code
batch_size = 8
in_channels = 16
out_channels = 64
kernel_size = (3, 5)
height_in = 64
width_in = 64

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization
# Convolution / Dense_Convolution_1D
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 1D convolution operation.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 1D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, length).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, length_out).
        """
        return self.conv1d(x)

# Test code
batch_size = 32
in_channels = 64
out_channels = 128
kernel_size = 3
length = 256

def get_inputs():
    x = torch.rand(batch_size, in_channels, length)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization
# Convolution / Dense_Convolution_2D
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs standard Dense (regular) 2D Convolution.
    All input channels are connected to all output channels.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0):
        """
        Initializes the Dense Conv2d layer.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            kernel_size (int): Size of the convolution kernel.
            stride (int): Stride of the convolution. Default: 1.
            padding (int): Padding added to input. Default: 0.
        """
        super(Model, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True
        )

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

        Args:
            x (torch.Tensor): Input tensor of shape (N, C_in, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (N, C_out, H_out, W_out).
        """
        return self.conv(x)

batch_size = 2
in_channels = 16
out_channels = 64
height = 64
width = 64
kernel_size = 3

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, 1, 1]  # stride=1, padding=1
# Convolution / Dense_Convolution_2D_kernel
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Standard 7x7 convolution, padding=3 to keep spatial size.
    """
    def __init__(self, in_ch: int = 3, out_ch: int = 32):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=7, padding=3, stride=1, bias=True)

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

batch_size = 2
in_channels = 3
out_channels = 32
height = 224
width = 224

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

def get_init_inputs():
    return [in_channels, out_channels]
# Convolution / Dense_Convolution_2D__asymmetric_input__square_kernel
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 2D convolution operation with an asymmetric input and a square kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, (kernel_size, kernel_size), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 2D convolution.

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

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv2d(x)

# Test code
batch_size = 8
# smaller spatial dims
height = 64
width = 128
in_channels = 32  # increased channels
out_channels = 64
kernel_size = 3
# asymmetric input: make width considerably larger than height

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization
# Convolution / Group_Convolution_2D
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Group Convolution.
    Input channels are divided into groups, each group is convolved independently.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, groups: int, 
                 stride: int = 1, padding: int = 0):
        """
        Initializes the Group Conv2d layer.

        Args:
            in_channels (int): Number of input channels (must be divisible by groups).
            out_channels (int): Number of output channels (must be divisible by groups).
            kernel_size (int): Size of the convolution kernel.
            groups (int): Number of groups for grouped convolution.
            stride (int): Stride of the convolution. Default: 1.
            padding (int): Padding added to input. Default: 0.
        """
        super(Model, self).__init__()
        assert in_channels % groups == 0, "in_channels must be divisible by groups"
        assert out_channels % groups == 0, "out_channels must be divisible by groups"
        
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=groups,
            bias=True
        )

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

        Args:
            x (torch.Tensor): Input tensor of shape (N, C_in, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (N, C_out, H_out, W_out).
        """
        return self.conv(x)

batch_size = 2
in_channels = 16
out_channels = 32
height = 64
width = 64
kernel_size = 3
groups = 4  # Split into 4 groups

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, groups, 1, 1]  # stride=1, padding=1