Jhnony/Mobilekernelbench huggingface dataset:
# Activation / Celu
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, alpha=1.0):
        super().__init__()
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.celu(x, alpha=self.alpha)

batch_size = 128
dim = 1024
alpha = 1.0

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

def get_init_inputs():
    return [alpha]
# Activation / ELU
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 the Exponential Linear Unit (ELU) activation function.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        ELU(x) = x if x > 0
        ELU(x) = alpha * (exp(x) - 1) if x <= 0
        Applies element-wise activation with configurable alpha parameter.
    """
    
    def __init__(self, alpha: float = 1.0):
        """
        Initialize the ELU model.
        
        Args:
            alpha: Scaling factor for negative values (default: 1.0)
        """
        super(Model, self).__init__()
        self.alpha = alpha
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor with ELU activation applied
        """
        return F.elu(x, alpha=self.alpha)


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

dim1 = 16
dim2 = 64
dim3 = 128
alpha = 1.0

def get_inputs():
    x = torch.randn(dim1, dim2, dim3)
    return [x]

def get_init_inputs():
    return [alpha]
# Activation / HardSigmoid
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a HardSigmoid activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies HardSigmoid activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with HardSigmoid applied, same shape as input.
        """
        return torch.nn.functional.hardsigmoid(x)

batch_size = 128
dim = 1024

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Activation / LogSoftmax
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a LogSoftmax activation.
    """
    def __init__(self, dim: int = 1):
        super(Model, self).__init__()
        self.dim = dim
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies LogSoftmax activation to the input tensor.

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

        Returns:
            torch.Tensor: Output tensor with LogSoftmax applied, same shape as input.
        """
        return torch.log_softmax(x, dim=self.dim)

batch_size = 32
dim = 512

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Activation / LogSoftmax_axis_0
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 LogSoftmax operation to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies log(softmax(x)) along the specified axis.
    """
    
    def __init__(self, axis: int = 0):
        super(Model, self).__init__()
        self.axis = axis
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying log_softmax
        """
        return F.log_softmax(x, dim=self.axis)


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

dim0 = 8
dim1 = 32
dim2 = 64

def get_inputs():
    x = torch.randn(dim0, dim1, dim2)
    return [x]

def get_init_inputs():
    return [0]  # [axis]
# Activation / LogSoftmax_axis_1
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 LogSoftmax operation to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies log(softmax(x)) along the specified axis.
    """
    
    def __init__(self, axis: int = 1):
        super(Model, self).__init__()
        self.axis = axis
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying log_softmax
        """
        return F.log_softmax(x, dim=self.axis)


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

dim0 = 8
dim1 = 32
dim2 = 64

def get_inputs():
    x = torch.randn(dim0, dim1, dim2)
    return [x]

def get_init_inputs():
    return [1]  # [axis]
# Activation / LogSoftmax_axis_2
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 LogSoftmax operation to the input tensor.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Applies log(softmax(x)) along the specified axis.
    """
    
    def __init__(self, axis: int = 2):
        super(Model, self).__init__()
        self.axis = axis
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying log_softmax
        """
        return F.log_softmax(x, dim=self.axis)


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

dim0 = 8
dim1 = 32
dim2 = 64

def get_inputs():
    x = torch.randn(dim0, dim1, dim2)
    return [x]

def get_init_inputs():
    return [2]  # [axis]
# Activation / PRelu
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_slopes):
        super().__init__()
        self.slope = nn.Parameter(torch.randn(num_slopes))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.prelu(x, self.slope)

batch_size = 64
dim = 1024

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

def get_init_inputs():
    return [dim]
# Activation / Relu
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a ReLU activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies ReLU activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with ReLU applied, same shape as input.
        """
        return torch.relu(x)

batch_size = 64
dim = 1024

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Activation / Softmax
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a Softmax activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Softmax activation to the input tensor.

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

        Returns:
            torch.Tensor: Output tensor with Softmax applied, same shape as input.
        """
        return torch.softmax(x, dim=1)

batch_size = 64
dim = 2048

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Activation / Softplus_2d
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a Softplus activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Softplus activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with Softplus applied, same shape as input.
        """
        return torch.nn.functional.softplus(x)

batch_size = 8
dim = 128

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Activation / Softplus_3d
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a Softplus activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Softplus activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with Softplus applied, same shape as input.
        """
        return torch.nn.functional.softplus(x)

batch_size = 8
dim1 = 128
dim2 = 128

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

def get_init_inputs():
    return []  # No special initialization inputs needed
# Binary / Add
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise addition (Add) of two tensors, with broadcasting.

    Semantics:
        y = x + y
    """

    def __init__(self):
        super(Model, self).__init__()

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return torch.add(a, b)


# ======== Example input configuration ========
# a: (batch, C, H, W)
# b: (1, C, 1, 1)
batch_size = 4
C, H, W = 3, 64, 64

def get_inputs():
    a = torch.rand(batch_size, C, H, W)
    b = torch.rand(1, C, 1, 1)
    return [a, b]

def get_init_inputs():
    return []
# Binary / Div
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise division (Div) with broadcasting.

    Semantics:
        y = a / b
    """

    def __init__(self):
        super(Model, self).__init__()

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return torch.div(a, b)


# ======== Example input configuration ========
# a: (batch, C, H, W)
# b: (1, C, 1, 1)
batch_size = 4
C, H, W = 3, 64, 64

def get_inputs():
    a = torch.rand(batch_size, C, H, W)  # numerator
    b = 0.5 + torch.rand(1, C, 1, 1)
    return [a, b]

def get_init_inputs():
    return []
# Binary / Min
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes element-wise minimum of two tensors.

    Semantics:
        y = min(x1, x2)
        Computes the element-wise minimum of two input tensors.
        Input shape: Any broadcastable shapes
        Output shape: Broadcasted shape of inputs
    """

    def __init__(self):
        super(Model, self).__init__()

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x1: First input tensor
            x2: Second input tensor
        
        Returns:
            Tensor containing element-wise minimum
        """
        return torch.min(x1, x2)


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

batch_size = 16
dim1 = 64
dim2 = 512

def get_inputs():
    x1 = torch.randn(batch_size, dim1, dim2)
    x2 = torch.randn(batch_size, dim1, dim2)
    return [x1, x2]

def get_init_inputs():
    return []
# Binary / Mod
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise modulo operation.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise modulo C = A mod B with broadcasting support.
        - If fmod=0: Uses sign of divisor (Python/NumPy style, torch.remainder)
        - If fmod=1: Uses sign of dividend (C/C++ style, torch.fmod)
        
        For fmod=0: C = A - B * floor(A/B)
        For fmod=1: C = A - B * trunc(A/B)
    """
    
    def __init__(self, fmod: int = 0):
        """
        Initialize the Mod model.
        
        Args:
            fmod: Modulo mode (0 for remainder, 1 for fmod)
        """
        super(Model, self).__init__()
        self.fmod = fmod
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: Dividend tensor
            b: Divisor tensor
        
        Returns:
            Element-wise modulo result with broadcasting
        """
        if self.fmod == 1:
            return torch.fmod(a, b)
        else:
            return torch.remainder(a, b)


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

dim1 = 8
dim2 = 256
dim3 = 128
fmod = 0  # 0 for remainder (Python style), 1 for fmod (C style)

def get_inputs():
    a = torch.randn(dim1, dim2, dim3)
    b = torch.randn(dim1, dim2, dim3)
    return [a, b]

def get_init_inputs():
    return [fmod]
# Binary / Mul
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise multiplication.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise multiplication C = A * B with broadcasting support.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the Mul model.
        """
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor
            b: Second input tensor
        
        Returns:
            Element-wise multiplication result with broadcasting
        """
        return torch.mul(a, b)


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

dim1 = 8 
dim2 = 256
dim3 = 256

def get_inputs():
    a = torch.randn(dim1, dim2, dim3)
    b = torch.randn(dim1, dim2, dim3)
    return [a, b]

def get_init_inputs():
    return []
# Binary / Mul_bcast
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise multiplication with broadcasting.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs element-wise multiplication: output = x * y
        y is broadcasted to match x's last dimension.
        x[8, 16, 16] * y[16] -> output[8, 16, 16]
    """
    
    def __init__(self):
        """
        Initialize the Mul model.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: First input tensor with shape [batch, dim1, dim2]
            y: Second input tensor with shape [dim2]
        
        Returns:
            Output tensor with element-wise multiplication applied
        """
        return x * y


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

batch_size = 8
dim1 = 64
dim2 = 64

def get_inputs():
    x = torch.randn(batch_size, dim1, dim2)
    y = torch.randn(dim2)  # 修改为 dim2=16 以匹配 x 的最后一个维度
    return [x, y]

def get_init_inputs():
    return []
# Binary / Pow
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise power operation.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise power C = A^B with broadcasting support.
        For each element: C[i] = A[i] ** B[i]
        Supports NumPy-style broadcasting for tensors of different shapes.
        
        Note: For negative bases with non-integer exponents, result will be NaN.
    """
    
    def __init__(self):
        """
        Initialize the Pow model.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Base tensor
            y: Exponent tensor
        
        Returns:
            Element-wise power result with broadcasting
        """
        return torch.pow(x, y)


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

dim1 = 8
dim2 = 256
dim3 = 256

def get_inputs():
    torch.manual_seed(42)
    x = torch.rand(dim1, dim2, dim3) * 10  # [0, 10) positive values
    y = torch.randn(dim1, dim2, dim3) * 2  # [-4, 4) any real numbers
    return [x, y]

def get_init_inputs():
    return []
# Binary / Sub
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise subtraction.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise subtraction C = A - B with broadcasting support.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the Sub model.
        """
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: Minuend tensor (被减数)
            b: Subtrahend tensor (减数)
        
        Returns:
            Element-wise subtraction result with broadcasting
        """
        return torch.sub(a, b)


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

dim1 = 8
dim2 = 128
dim3 = 256

def get_inputs():
    a = torch.randn(dim1, dim2, dim3)
    b = torch.randn(dim1, dim2, dim3)
    return [a, b]

def get_init_inputs():
    return []
# Convolution / ConvolutionDepthwise
import torch
import torch.nn as nn

class DepthwiseSeparableConv2d(nn.Module):
    """
    Depthwise Separable Convolution as used in MobileNet.
    Consists of: Depthwise Conv (groups=in_channels) + Pointwise Conv (1x1).
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0):
        """
        Initializes Depthwise Separable Convolution.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            kernel_size (int): Size of the depthwise convolution kernel.
            stride (int): Stride of the depthwise convolution. Default: 1.
            padding (int): Padding for depthwise convolution. Default: 0.
        """
        super(DepthwiseSeparableConv2d, self).__init__()
        
        # Depthwise convolution: each channel is convolved independently
        self.depthwise = nn.Conv2d(
            in_channels=in_channels,
            out_channels=in_channels,  # Same as input channels
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=in_channels,  # Key: groups equals in_channels
            bias=False
        )
        
        # Pointwise convolution: 1x1 conv to mix channels
        self.pointwise = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=1,
            bias=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Depthwise Separable Convolution.

        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).
        """
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

batch_size = 16
in_channels = 16
out_channels = 32
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 / ConvTranspose_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 = 16
in_channels = 16
out_channels = 32
kernel_size = (3, 5)
height_in = 32
width_in = 32

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