Jhnony/Mobilekernelbench huggingface dataset:
# Matrix / Einsum_elementwise
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise multiplication using Einsum.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs element-wise multiplication of two matrices
        Equation: "ij,ij->ij"
    """
    
    def __init__(self, equation: str = "ij,ij->ij"):
        super(Model, self).__init__()
        self.equation = equation
    
    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x1: First input tensor with shape [M, N]
            x2: Second input tensor with shape [M, N]
        
        Returns:
            Output tensor with shape [M, N]
        """
        return torch.einsum(self.equation, x1, x2)


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

M = 256
N = 256
equation = "ij,ij->ij"

def get_inputs():
    x1 = torch.randn(M, N)
    x2 = torch.randn(M, N)
    return [x1, x2]

def get_init_inputs():
    return [equation]
# Matrix / Einsum_sum_all
import torch
import torch.nn as nn


class Model(nn.Module):
    """
    A model that sums all elements using Einsum.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Sums all elements in the tensor
        Equation: "ij->"
    """
    
    def __init__(self, equation: str = "ij->"):
        super(Model, self).__init__()
        self.equation = equation
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [M, N]
        
        Returns:
            Scalar tensor containing sum of all elements
        """
        return torch.einsum(self.equation, x)


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

M = 64
N = 256
equation = "ij->"

def get_inputs():
    x = torch.randn(M, N)
    return [x]

def get_init_inputs():
    return [equation]
# Matrix / Einsum_transpose
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs transpose using Einsum.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Transposes a 2D matrix
        Equation: "ij->ji"
    """
    
    def __init__(self, equation: str = "ij->ji"):
        super(Model, self).__init__()
        self.equation = equation
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [M, N]
        
        Returns:
            Output tensor with shape [N, M]
        """
        return torch.einsum(self.equation, x)


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

M = 64
N = 256
equation = "ij->ji"

def get_inputs():
    x = torch.randn(M, N)
    return [x]

def get_init_inputs():
    return [equation]
# Matrix / Gemm
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Basic GEMM (General Matrix Multiplication)
    Implements: Y = alpha * A @ B (+ beta * C)
    """
    def __init__(self, in_features: int, out_features: int, alpha: float = 1.0, beta: float = 1.0, with_bias: bool = True):
        super(Model, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.with_bias = with_bias
        self.linear = nn.Linear(in_features, out_features, bias=with_bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Y = alpha * (A @ W^T + b * beta)
        y = self.alpha * self.linear(x)
        return y

# ------------------- Hyperparameters -------------------
batch_size = 32
in_features = 512
out_features = 256
alpha = 1.0
beta = 1.0
with_bias = True

# ------------------- Input definitions -----------------
def get_inputs():
    x = torch.rand(batch_size, in_features)
    return [x]

def get_init_inputs():
    return [in_features, out_features, alpha, beta, with_bias]
# Matrix / Gemm_alpha
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs GEMM (General Matrix Multiplication) operation.
    
    Number of inputs: 3
    Implementation type: direct
    
    Semantics:
        Y = alpha * (A @ B) + beta * C
        Supports optional transpose of A and B matrices.
    """
    
    def __init__(self, alpha: float = 0.5, beta: float = 1.0, transA: int = 0, transB: int = 0):
        super(Model, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.transA = transA
        self.transB = transB
    
    def forward(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor with shape [M, K] or [K, M] if transA=1
            b: Second input tensor with shape [K, N] or [N, K] if transB=1
            c: Bias tensor that broadcasts to output shape
        
        Returns:
            Output tensor with shape [M, N]
        """
        a_mat = a.T if self.transA == 1 else a
        b_mat = b.T if self.transB == 1 else b
        return self.alpha * torch.mm(a_mat, b_mat) + self.beta * c


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

M = 8
K = 64
N = 128
alpha = 0.5
beta = 1.0
transA = 0
transB = 0

def get_inputs():
    a = torch.randn(M, K)
    b = torch.randn(K, N)
    c = torch.randn(1, N)
    return [a, b, c]

def get_init_inputs():
    return [alpha, beta, transA, transB]
# Matrix / Gemm_beta
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs GEMM (General Matrix Multiplication) operation.
    
    Number of inputs: 3
    Implementation type: direct
    
    Semantics:
        Y = alpha * (A @ B) + beta * C
        Supports optional transpose of A and B matrices.
    """
    
    def __init__(self, alpha: float = 0.5, beta: float = 1.0, transA: int = 0, transB: int = 0):
        super(Model, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.transA = transA
        self.transB = transB
    
    def forward(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor with shape [M, K] or [K, M] if transA=1
            b: Second input tensor with shape [K, N] or [N, K] if transB=1
            c: Bias tensor that broadcasts to output shape
        
        Returns:
            Output tensor with shape [M, N]
        """
        a_mat = a.T if self.transA == 1 else a
        b_mat = b.T if self.transB == 1 else b
        return self.alpha * torch.mm(a_mat, b_mat) + self.beta * c


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

M = 8
K = 64
N = 64
alpha = 1
beta = 0.5
transA = 0
transB = 0

def get_inputs():
    a = torch.randn(M, K)
    b = torch.randn(K, N)
    c = torch.randn(1, N)
    return [a, b, c]

def get_init_inputs():
    return [alpha, beta, transA, transB]
# Matrix / Gemm_no_bias
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs GEMM (General Matrix Multiplication) operation without bias.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Y = alpha * (A @ B)
        Matrix multiplication with optional transpose and scaling, no bias term.
    """
    
    def __init__(self, alpha: float = 1.0, transA: int = 0, transB: int = 0):
        super(Model, self).__init__()
        self.alpha = alpha
        self.transA = transA
        self.transB = transB
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor with shape [M, K] or [K, M] if transA=1
            b: Second input tensor with shape [K, N] or [N, K] if transB=1
        
        Returns:
            Output tensor with shape [M, N]
        """
        a_mat = a.T if self.transA == 1 else a
        b_mat = b.T if self.transB == 1 else b
        return self.alpha * torch.mm(a_mat, b_mat)


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

M = 16
K = 128
N = 32
alpha = 1.0
transA = 0
transB = 0

def get_inputs():
    a = torch.randn(M, K)
    b = torch.randn(K, N)
    return [a, b]

def get_init_inputs():
    return [alpha, transA, transB]
# Matrix / Gemm_scalar_bias
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Basic GEMM (General Matrix Multiplication)
    Implements: Y = alpha * A @ B (+ beta * C)
    """
    def __init__(self, in_features: int, out_features: int, alpha: float = 1.0, beta: float = 1.0, with_bias: bool = True):
        super(Model, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.with_bias = with_bias
        self.linear = nn.Linear(in_features, out_features, bias=with_bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Y = alpha * (A @ W^T + b * beta)
        y = self.alpha * self.linear(x)
        return y

# ------------------- Hyperparameters -------------------
batch_size = 32
in_features = 256
out_features = 512
alpha = 1.0
beta = 1.0
with_bias = True

# ------------------- Input definitions -----------------
def get_inputs():
    x = torch.rand(batch_size, in_features)
    return [x]

def get_init_inputs():
    return [in_features, out_features, alpha, beta, with_bias]
# Matrix / Gemm_transposed
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs GEMM with transposed input.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Y = alpha * (A^T @ W^T + b)
        Input is transposed before matrix multiplication.
    """
    
    def __init__(self, in_features: int, out_features: int, alpha: float = 1.0):
        super(Model, self).__init__()
        self.alpha = alpha
        self.linear = nn.Linear(in_features, out_features, bias=True)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [in_features, batch_size]
        
        Returns:
            Output tensor with shape [batch_size, out_features]
        """
        # Transpose input before linear transformation
        x_t = x.t()
        return self.alpha * self.linear(x_t)


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

batch_size = 32
in_features = 128
out_features = 256
alpha = 1.5

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

def get_init_inputs():
    return [in_features, out_features, alpha]
# Matrix / MatMul
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a single square matrix multiplication (C = A * B)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs the matrix multiplication.

        Args:
            A (torch.Tensor): Input matrix A of shape (N, N).
            B (torch.Tensor): Input matrix B of shape (N, N).

        Returns:
            torch.Tensor: Output matrix C of shape (N, N).
        """
        return torch.matmul(A, B)

N = 256
M = 128

def get_inputs():
    A = torch.rand(N, M)
    B = torch.rand(M, N)
    return [A, B]

def get_init_inputs():
    return []  # No special initialization inputs needed
# Matrix / MatMul_1d_3d
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs matrix multiplication on two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs matrix multiplication with broadcasting: output = a @ b
        When a is 1D and b is 3D, a is treated as a row vector.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor (1D vector)
            b: Second input tensor (multi-dimensional)
        
        Returns:
            Output tensor after matrix multiplication
        """
        return torch.matmul(a, b)


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

m = 16
batch_size = 32
p = 64

def get_inputs():
    a = torch.randn(m)
    b = torch.randn(batch_size, m, p)
    return [a, b]

def get_init_inputs():
    return []
# Matrix / MatMul_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 logical XOR operation to two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes the logical XOR of two boolean tensors element-wise.
        Supports NumPy-style broadcasting.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: First input tensor (boolean type)
            y: Second input tensor (boolean type)
        
        Returns:
            Output tensor after logical XOR operation
        """
        return torch.logical_xor(x, y)


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

dim0 = 32
dim1 = 256
dim2 = 256

def get_inputs():
    x = torch.randint(0, 2, (dim0, dim1, dim2)).bool()
    y = torch.randint(0, 2, (dim0, dim1, dim2)).bool()
    return [x, y]

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


class Model(nn.Module):
    """
    A model that performs matrix multiplication on two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs batched matrix multiplication: output = a @ b
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor with shape [batch_size, ..., n, m]
            b: Second input tensor with shape [batch_size, ..., m, p]
        
        Returns:
            Output tensor with shape [batch_size, ..., n, p]
        """
        return torch.matmul(a, b)


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

batch_dim1 = 8
batch_dim2 = 1
n = 16
m = 32
batch_dim3 = 1
batch_dim4 = 8
p = 32

def get_inputs():
    a = torch.randn(batch_dim1, batch_dim2, n, m)
    b = torch.randn(batch_dim3, batch_dim4, m, p)
    return [a, b]

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


class Model(nn.Module):
    """
    A model that performs matrix multiplication on two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs matrix multiplication: output = a @ b
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor with shape [n, m]
            b: Second input tensor with shape [m, p]
        
        Returns:
            Output tensor with shape [n, p]
        """
        return torch.matmul(a, b)


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

n = 128
m = 128
p = 128

def get_inputs():
    a = torch.randn(n, m)
    b = torch.randn(m, p)
    return [a, b]

def get_init_inputs():
    return []
# Normalization / BatchNormalization
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Batch Normalization.
    """
    def __init__(self, num_features: int):
        """
        Initializes the BatchNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
        """
        super(Model, self).__init__()
        self.bn = nn.BatchNorm2d(num_features=num_features)

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

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

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

batch_size = 16
features = 32
dim1 = 128
dim2 = 128

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

def get_init_inputs():
    return [features]
# Normalization / BatchNormalization_epsilon
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs Batch Normalization with custom epsilon.
    """
    def __init__(self, num_features: int, eps: float = 1e-5):
        """
        Initializes the BatchNorm layer with custom epsilon.

        Args:
            num_features (int): Number of features in the input tensor.
            eps (float, optional): Small value added to denominator for numerical stability. 
                                   Defaults to 1e-5.
        """
        super(Model, self).__init__()
        self.bn = nn.BatchNorm2d(num_features=num_features, eps=eps)

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

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

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

batch_size = 16
features = 32
dim1 = 128
dim2 = 128
epsilon = 1e-3  # Custom epsilon value (larger than default 1e-5)

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

def get_init_inputs():
    return [features, epsilon]
# Normalization / BatchNormalization_momentum
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs Batch Normalization with custom momentum.
    """
    def __init__(self, num_features: int, momentum: float = 0.1):
        """
        Initializes the BatchNorm layer with custom momentum.

        Args:
            num_features (int): Number of features in the input tensor.
            momentum (float, optional): Momentum for running mean and variance computation.
                                        Formula: running_stat = (1 - momentum) * running_stat + momentum * batch_stat
                                        Defaults to 0.1.
        """
        super(Model, self).__init__()
        self.bn = nn.BatchNorm2d(num_features=num_features, momentum=momentum)

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

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

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

batch_size = 16
features = 32
dim1 = 128
dim2 = 128
momentum = 0.2  # Custom momentum value (higher than default 0.1)

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

def get_init_inputs():
    return [features, momentum]
# Normalization / InstanceNormalization
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Instance Normalization.
    """
    def __init__(self, num_features: int):
        """
        Initializes the InstanceNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
        """
        super(Model, self).__init__()
        self.inorm = nn.InstanceNorm2d(num_features=num_features)

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

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

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

batch_size = 16
features = 32
dim1 = 64
dim2 = 64

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

def get_init_inputs():
    return [features]
# Normalization / InstanceNormalization_1d
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that applies 1D Instance Normalization.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes input over temporal/sequence dimension for each channel independently.
        Commonly used for sequence data or 1D signals.
    """
    
    def __init__(self, num_features: int, eps: float = 1e-5, affine: bool = False):
        """
        Args:
            num_features: Number of channels (C) in the input
            eps: Small value added to variance for numerical stability
            affine: If True, learnable affine parameters (gamma, beta)
        """
        super(Model, self).__init__()
        self.inorm = nn.InstanceNorm1d(
            num_features=num_features,
            eps=eps,
            affine=affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, num_features, length]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.inorm(x)


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

batch_size = 8
num_features = 32
length = 128
eps = 1e-5
affine = False

def get_inputs():
    x = torch.randn(batch_size, num_features, length)
    return [x]

def get_init_inputs():
    return [num_features, eps, affine]
# Normalization / InstanceNormalization_3d
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that applies 3D Instance Normalization.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Normalizes input over spatial dimensions (D, H, W) for each channel independently.
        Commonly used for 3D images (e.g., medical imaging, video).
    """
    
    def __init__(self, num_features: int, eps: float = 1e-5, affine: bool = False):
        """
        Args:
            num_features: Number of channels (C) in the input
            eps: Small value added to variance for numerical stability
            affine: If True, learnable affine parameters (gamma, beta)
        """
        super(Model, self).__init__()
        self.inorm = nn.InstanceNorm3d(
            num_features=num_features,
            eps=eps,
            affine=affine
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, num_features, depth, height, width]
        
        Returns:
            Output tensor with same shape as input
        """
        return self.inorm(x)


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

batch_size = 4
num_features = 16
depth = 8
height = 64
width = 64
eps = 1e-5
affine = False

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

def get_init_inputs():
    return [num_features, eps, affine]