Jhnony/Mobilekernelbench huggingface dataset:
# Convolution / Group_Convolution_2D_kernel
import torch
import torch.nn as nn

class GroupedConvDepthwiseLike(nn.Module):
    """
    Grouped convolution with groups equal to half of input channels.
    """
    def __init__(self, in_ch: int = 24, out_ch: int = 24,kernel_size: int =5, groups: int =12, stride: int =1, padding: int =2):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=kernel_size, padding=padding, stride=stride, groups=groups, bias=False)

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

batch_size = 2
def get_inputs():
    x = torch.rand(batch_size, 24, 32, 32)
    return [x]

def get_init_inputs():
    return []


batch_size = 2
in_channels = 24
out_channels = 24
height = 64
width = 64
kernel_size = 5
groups = 12  # Split into 12 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, 2]  # stride=1, padding=2
# Convolution / Strassen_Convolution_2D
import torch
import torch.nn as nn

class StrassenConv1x1Stride2(nn.Module):
    """
    1x1 convolution (conceptually could use Strassen for the GEMM part), stride=2 for downsampling.
    """
    def __init__(self, in_ch: int = 64, out_ch: int = 128):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=1, stride=2, padding=0, bias=True)

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

batch_size = 4
in_channels = 64
out_channels = 128
height = 56
width = 56

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

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

class Model(nn.Module):
    """
    Model using 3x3 convolutions which are optimized by Winograd algorithm in cuDNN.
    Winograd reduces multiplications for small kernels (especially 3x3).
    
    For F(2x2, 3x3): reduces from 36 to 16 multiplications.
    cuDNN automatically selects Winograd when beneficial.
    """
    def __init__(self, in_channels: int, out_channels: int, use_winograd_hint: bool = True):
        """
        Initializes Conv2d layer optimized for Winograd algorithm.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            use_winograd_hint (bool): If True, use 3x3 kernel (Winograd-friendly).
        """
        super(Model, self).__init__()
        
        # 3x3 convolution is optimal for Winograd
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,  # Winograd is most efficient for 3x3
            stride=1,       # Stride=1 for Winograd compatibility
            padding=1,
            bias=False
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Winograd-optimized convolution (handled by cuDNN backend).

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

        Returns:
            torch.Tensor: Output tensor of shape (N, C_out, H, W).
        """
        x = self.conv(x)
        return x

batch_size = 8
in_channels = 16
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, True]
# Convolution / Winograd_Convolution_2D_padding
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    3x3 convolution (may be accelerated by Winograd in backend), dilation=2.
    """
    def __init__(self, in_ch: int = 16, out_ch: int = 16):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=3, padding=2, dilation=2, stride=1, bias=True)

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

batch_size = 4
in_channels = 16
out_channels = 16
height = 64
width = 64

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

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

class Model(nn.Module):
    """
    Element-wise logical AND between two boolean tensors.

    Semantics:
        y = logical_and(a, b)
    """

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

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        a_bool = a.bool()
        b_bool = b.bool()
        return torch.logical_and(a_bool, b_bool)


# ======== Example input configuration ========
batch_size = 4
H, W = 32, 64

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

def get_init_inputs():
    return []
# Logic / BitwiseAnd
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 bitwise AND operation to two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes the bitwise AND of two integer 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 (integer type)
            y: Second input tensor (integer type)
        
        Returns:
            Output tensor after bitwise AND operation
        """
        return torch.bitwise_and(x, y)


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

dim0 = 8
dim1 = 32
dim2 = 64

def get_inputs():
    x = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    y = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    return [x, y]

def get_init_inputs():
    return []
# Logic / BitWiseOr
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 bitwise OR operation to two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes the bitwise OR of two integer 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 (integer type)
            y: Second input tensor (integer type)
        
        Returns:
            Output tensor after bitwise OR operation
        """
        return torch.bitwise_or(x, y)


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

dim0 = 8
dim1 = 128
dim2 = 256

def get_inputs():
    x = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    y = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    return [x, y]

def get_init_inputs():
    return []
# Logic / BitWiseXor
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 bitwise XOR operation to two input tensors.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes the bitwise XOR of two integer 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 (integer type)
            y: Second input tensor (integer type)
        
        Returns:
            Output tensor after bitwise XOR operation
        """
        return torch.bitwise_xor(x, y)


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

dim0 = 8
dim1 = 128
dim2 = 512

def get_inputs():
    x = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    y = torch.randint(0, 256, (dim0, dim1, dim2), dtype=torch.int32)
    return [x, y]

def get_init_inputs():
    return []
# Logic / Equal
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise equality comparison.

    Semantics:
        y = (a == b)
    Output:
        bool tensor, broadcast semantics like ONNX Equal.
    """

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

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


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

batch_size = 4
C, H, W = 3, 64, 128

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

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


class Model(nn.Module):
    """
    A model that performs element-wise greater-than comparison.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise comparison C = (A > B) with broadcasting support.
        Returns boolean tensor where each element is True if A > B, False otherwise.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the Greater model.
        
        Note: This operator has no learnable parameters.
        """
        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:
            Boolean tensor with element-wise comparison result (a > b)
        """
        return torch.gt(a, b)


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

dim1 = 8
dim2 = 128
dim3 = 512

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

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


class Model(nn.Module):
    """
    A model that performs element-wise greater-than-or-equal comparison.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise comparison C = (A >= B) with broadcasting support.
        Returns boolean tensor where each element is True if A >= B, False otherwise.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the GreaterOrEqual model.
        
        Note: This operator has no learnable parameters.
        """
        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:
            Boolean tensor with element-wise comparison result (a >= b)
        """
        return torch.ge(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 []
# Logic / Less
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise less-than comparison.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise comparison C = (A < B) with broadcasting support.
        Returns boolean tensor where each element is True if A < B, False otherwise.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the Less model.
        
        Note: This operator has no learnable parameters.
        """
        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:
            Boolean tensor with element-wise comparison result (a < b)
        """
        return torch.lt(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 []
# Logic / LessEqual
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise less-than-or-equal comparison.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise comparison C = (A <= B) with broadcasting support.
        Returns boolean tensor where each element is True if A <= B, False otherwise.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the LessOrEqual model.
        
        Note: This operator has no learnable parameters.
        """
        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:
            Boolean tensor with element-wise comparison result (a <= b)
        """
        return torch.le(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 []
# Logic / LogicalXor
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs element-wise logical XOR operation.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Computes element-wise logical XOR: C = A XOR B with broadcasting support.
        Returns boolean tensor where each element is True if exactly one of A or B is True.
        Non-zero values are treated as True, zero values as False.
        Supports NumPy-style broadcasting for tensors of different shapes.
    """
    
    def __init__(self):
        """
        Initialize the Xor model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            a: First input tensor (numeric or boolean)
            b: Second input tensor (numeric or boolean)
        
        Returns:
            Boolean tensor with element-wise logical XOR result
        """
        return torch.logical_xor(a, b)


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

dim1 = 8
dim2 = 128
dim3 = 128

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

def get_init_inputs():
    return []
# Logic / LogicalXor3d
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 = 16
dim1 = 128
dim2 = 64

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 []
# Logic / LogicalXor_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 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 = 64
dim1 = 512

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

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


class Model(nn.Module):
    """
    A model that performs element-wise logical NOT operation.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes element-wise logical negation of the input tensor.
        Non-zero values are treated as True, zero values as False.
        Output is boolean tensor where each element is the negation of input.
        
        Y[i] = NOT X[i]
    """
    
    def __init__(self):
        """
        Initialize the Not model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor of any numeric or boolean type
        
        Returns:
            Boolean tensor with element-wise logical negation
        """
        return torch.logical_not(x)


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

dim1 = 8
dim2 = 64
dim3 = 64

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

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

class Model(nn.Module):
    """
    A model that computes the determinant (Det) of the input tensor.

    Semantics:
        y = det(x)
        Computes the determinant of each square matrix in the input tensor.
        Input shape: (..., n, n) where the last two dimensions form square matrices.
        Output shape: (...) with determinant values for each matrix.
    """

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input tensor of shape (..., n, n) containing square matrices.
        
        Returns:
            Tensor of shape (...) containing determinant of each matrix.
        """
        return torch.linalg.det(x)


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

matrix_size = 512  # n x n square matrix
input_shape = (matrix_size, matrix_size)

def get_inputs():
    # Generate random square matrices
    x = 2.0 * torch.rand(*input_shape) - 1.0
    return [x]

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


class Model(nn.Module):
    """
    A model that computes the determinant of 2x2 matrices.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes the determinant of each 2x2 matrix in the input tensor.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [..., 2, 2]
        
        Returns:
            Output tensor with shape [...] containing determinants
        """
        return torch.linalg.det(x)


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

batch_size = 64
matrix_size = 512

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

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


class Model(nn.Module):
    """
    A model that computes the determinant with multiple batch dimensions.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes the determinant of matrices with multiple leading batch dimensions.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch1, batch2, n, n]
        
        Returns:
            Output tensor with shape [batch1, batch2] containing determinants
        """
        return torch.linalg.det(x)


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

batch_size1 = 4
batch_size2 = 128
matrix_size = 64

def get_inputs():
    x = torch.randn(batch_size1, batch_size2, matrix_size, matrix_size)
    return [x]

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


class Model(nn.Module):
    """
    A model that performs matrix multiplication using Einsum.
    
    Number of inputs: 2
    Implementation type: direct
    
    Semantics:
        Performs standard matrix multiplication: C = A @ B
        Equation: "ij,jk->ik"
    """
    
    def __init__(self, equation: str = "ij,jk->ik"):
        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, K]
            x2: Second input tensor with shape [K, N]
        
        Returns:
            Output tensor with shape [M, N]
        """
        return torch.einsum(self.equation, x1, x2)


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

M = 32
K = 64
N = 128
equation = "ij,jk->ik"

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

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


class Model(nn.Module):
    """
    A model that extracts diagonal elements using Einsum.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Extracts the diagonal of a square matrix
        Equation: "ii->i"
    """
    
    def __init__(self, equation: str = "ii->i"):
        super(Model, self).__init__()
        self.equation = equation
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [N, N]
        
        Returns:
            Output tensor with shape [N] containing diagonal elements
        """
        return torch.einsum(self.equation, x)


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

N = 1024
equation = "ii->i"

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

def get_init_inputs():
    return [equation]