Jhnony/Mobilekernelbench huggingface dataset:
# Tensor / CumSum
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that applies cumulative sum (CumSum) to the input tensor.

    Semantics:
        y = cumsum(x, dim=axis)
        Computes the cumulative sum of elements along the specified dimension.
    """

    def __init__(self, axis: int = -1):
        """
        Args:
            axis: The dimension along which to compute cumulative sum. Default is -1 (last dimension).
        """
        super(Model, self).__init__()
        self.axis = axis

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.cumsum(x, dim=self.axis)


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

batch_size = 16
input_shape = (8, 32, 32)

def get_inputs():
    x = 2.0 * torch.rand(batch_size, *input_shape) - 1.0
    return [x]

def get_init_inputs():
    return []
# Tensor / DepthToSpace
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a depth-to-space (pixel shuffle) operation.

    This is equivalent to the DepthToSpace operator with a given block size
    in NCHW layout. Input shape is (N, C * block_size^2, H, W) and the output
    shape is (N, C, H * block_size, W * block_size).

    Parameters:
        block_size (int): The spatial upscaling factor (DepthToSpace blockSize).
    """

    def __init__(self, block_size: int):
        super(Model, self).__init__()
        self.block_size = block_size
        # In ONNX export, this will map to a DepthToSpace op (mode = DCR).
        self.op = nn.PixelShuffle(block_size)

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


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

# DepthToSpace: (N, C_in, H_in, W_in)
# C_in = C_out * block_size^2
batch_size = 4
block_size = 4

out_channels = 8       # C_out
out_height = 64        # H_out
out_width = 64         # W_out

in_channels = out_channels * (block_size ** 2)  # C_in = C_out * r^2
in_height = out_height // block_size            # H_in = H_out / r
in_width = out_width // block_size              # W_in = W_out / r

input_shape = (in_channels, in_height, in_width)


def get_inputs():
    # NCHW: (batch_size, C_in, H_in, W_in)
    return [torch.rand(batch_size, *input_shape)]


def get_init_inputs():
    # Constructor args for Model(...)
    return [block_size]
# Tensor / Gather
# Gather.py
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a Gather operation along a given axis.

    Semantics:
        y = gather(params, axis, indices)

    Here we implement the common case where `indices` is 1D:
        params: Tensor of shape (d0, d1, ..., d_{k-1})
        axis:   integer in [0, k-1]
        indices: 1D Long tensor of shape (n,)

    Output shape:
        (d0, ..., d_{axis-1}, n, d_{axis+1}, ..., d_{k-1})
    """

    def __init__(self, axis: int):
        super(Model, self).__init__()
        self.axis = axis

    def forward(self, params: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
        # indices must be 1D Long tensor
        assert indices.dim() == 1, "This standard case expects 1D indices"
        assert indices.dtype in (torch.int32, torch.int64)
        if indices.dtype != torch.int64:
            indices = indices.to(torch.int64)

        axis = self.axis if self.axis >= 0 else self.axis + params.dim()
        return torch.index_select(params, dim=axis, index=indices)


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

# params shape: (B, C, D)
B, C, D = 8, 64, 64
axis = 1  # gather along C-dimension

def get_inputs():
    params = torch.rand(B, C, D)
    # indices in [0, C)
    idx = torch.randint(low=0, high=C, size=(3,), dtype=torch.int64)
    return [params, idx]

def get_init_inputs():
    # Constructor arguments for Model(...)
    return [axis]
# Tensor / GatherElements
# GatherElements.py
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a GatherElements operation along a given axis.

    Semantics (same as ONNX GatherElements / torch.gather):
        output = gather_elements(data, axis, indices)

    Requirements:
        - data and indices have the same shape
        - indices contains valid indices along `axis`
    """

    def __init__(self, axis: int):
        super(Model, self).__init__()
        self.axis = axis

    def forward(self, data: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
        assert data.shape == indices.shape, "data and indices must have the same shape"
        if indices.dtype != torch.int64:
            indices = indices.to(torch.int64)

        axis = self.axis if self.axis >= 0 else self.axis + data.dim()
        return torch.gather(data, dim=axis, index=indices)


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

# data shape: (B, C, D)
B, C, D = 16, 64, 64
axis = 1  # gather along C-dimension

def get_inputs():
    data = torch.rand(B, C, D)
    # indices must be in [0, C) for axis=1
    idx = torch.randint(low=0, high=C, size=(B, C, D), dtype=torch.int64)
    return [data, idx]

def get_init_inputs():
    return [axis]
# Tensor / GridSample
import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    GridSample: sample input using a normalized grid.

    Semantics (PyTorch / ONNX-like):
        y = grid_sample(input, grid, mode, padding_mode, align_corners)

    Here:
        input: (N, C, H, W)
        grid:  (N, H_out, W_out, 2), values in [-1, 1]
    """

    def __init__(self, mode: str = "bilinear",
                       padding_mode: str = "zeros",
                       align_corners: bool = False):
        super(Model, self).__init__()
        self.mode = mode
        self.padding_mode = padding_mode
        self.align_corners = align_corners

    def forward(self, x: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
        return F.grid_sample(
            x, grid,
            mode=self.mode,
            padding_mode=self.padding_mode,
            align_corners=self.align_corners,
        )


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

N, C, H, W = 4, 3, 32, 64
H_out, W_out = 4, 4

def get_inputs():
    x = torch.rand(N, C, H, W)
    # grid in [-1, 1]
    grid = 2.0 * torch.rand(N, H_out, W_out, 2) - 1.0
    return [x, grid]

def get_init_inputs():
    # 默认 bilinear / zeros / align_corners=False
    return ["bilinear", "zeros", False]
# Tensor / OneHot
import torch
import torch.nn.functional as F
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, depth, on_value=1, off_value=0, axis=-1):
        super(Model, self).__init__()
        self.depth = depth
        self.on_value = on_value
        self.off_value = off_value
        self.axis = axis

    def forward(self, indices: torch.Tensor) -> torch.Tensor:
        one_hot_result = F.one_hot(indices, num_classes=self.depth)
        one_hot_result = one_hot_result * (self.on_value - self.off_value) + self.off_value

        if self.axis != -1:
            one_hot_result = one_hot_result.permute(*range(one_hot_result.dim() - 1), self.axis, -1)
        
        return one_hot_result

batch_size = 16
dim1 = 128
dim2 = 1024
depth = 10

# ======== Example input configuration ========
def get_inputs() -> list:
    x = torch.randint(0, depth, (batch_size, dim1))
    return [x]

def get_init_inputs():
    return [depth]
# Tensor / Range
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, start, limit, delta):
        return torch.arange(start, limit, delta)

start = torch.tensor(0.0)
limit = torch.tensor(1024.0)
delta = torch.tensor(1.0)

# ======== Example input configuration ========
def get_inputs() -> list:
    return [start, limit, delta]

def get_init_inputs():
    return []
# Tensor / Reshape
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a reshape operation on the input tensor.

    The input is assumed to be of shape (batch_size, *input_shape),
    and it will be reshaped to `target_shape`.

    Parameters:
        target_shape (tuple or list of int): The target shape passed to
            torch.reshape. It must be compatible with the input number of
            elements. You can use -1 for one dimension to let PyTorch infer it.
    """

    def __init__(self, target_shape):
        super(Model, self).__init__()
        self.target_shape = tuple(target_shape)

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


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

batch_size = 16
input_shape = (8, 64, 64)  # C, H, W

target_shape = (batch_size, -1)

def get_inputs():
    return [torch.rand(batch_size, *input_shape)]

def get_init_inputs():
    return [target_shape]
# Tensor / ScatterElements
# ScatterElements.py
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a ScatterElements-like operation along a given axis.

    Semantics (assign mode):
        out = data.clone()
        out.scatter_(dim=axis, index=indices, src=updates)

    Requirements:
        - data, indices, updates have the same shape
        - indices contains valid indices along `axis`
    """

    def __init__(self, axis: int):
        super(Model, self).__init__()
        self.axis = axis

    def forward(self, data: torch.Tensor,
                      indices: torch.Tensor,
                      updates: torch.Tensor) -> torch.Tensor:
        assert data.shape == indices.shape == updates.shape, \
            "data, indices and updates must have the same shape"

        if indices.dtype != torch.int64:
            indices = indices.to(torch.int64)

        axis = self.axis if self.axis >= 0 else self.axis + data.dim()

        out = data.clone()
        out.scatter_(dim=axis, index=indices, src=updates)
        return out


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

B, C, L = 8, 64, 64
data_shape = (B, C, L)
axis = 1

def get_inputs():
    data = torch.zeros(*data_shape)
    indices = torch.randint(low=0, high=C, size=data_shape, dtype=torch.int64)
    updates = torch.rand(*data_shape)

    return [data, indices, updates]

def get_init_inputs():
    return [axis]
# Tensor / Slice
# Slice.py
import torch
import torch.nn as nn


class Model(nn.Module):
    """
    ONNX Slice-equivalent model.
    Extracts a single slice from the input tensor.
    """

    def __init__(self, axis: int, start: int, end: int):
        super().__init__()
        self.axis = axis
        self.start = start
        self.end = end

    def forward(self, x: torch.Tensor):
        axis = self.axis if self.axis >= 0 else self.axis + x.dim()

        # Build slicing objects for all dimensions
        slices = [slice(None)] * x.dim()
        slices[axis] = slice(self.start, self.end)

        return x[tuple(slices)]

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

batch_size = 2
input_shape = (10, 4, 4)   # C, H, W

# Slice configuration (ONNX Slice semantics)
axis = 1
start = 3
end = 6   # end is exclusive, matches ONNX Slice

def get_inputs():
    # Input shape: (N, C, H, W) = (2, 10, 4, 4)
    x = torch.rand(batch_size, *input_shape)
    return [x]

def get_init_inputs():
    # Model __init__(axis, start, end)
    return [axis, start, end]
# Tensor / StridedSlice
# Slice.py
import torch
import torch.nn as nn


class Model(nn.Module):
    """
    ONNX Slice-equivalent model.
    Extracts a single slice from the input tensor.
    """

    def __init__(self, axis: int, start: int, end: int):
        super().__init__()
        self.axis = axis
        self.start = start
        self.end = end

    def forward(self, x: torch.Tensor):
        axis = self.axis if self.axis >= 0 else self.axis + x.dim()

        # Build slicing objects for all dimensions
        slices = [slice(None)] * x.dim()
        slices[axis] = slice(self.start, self.end)

        return x[tuple(slices)]

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

batch_size = 32
input_shape = (8, 64, 64)   # C, H, W

# Slice configuration (ONNX Slice semantics)
axis = 1
start = 3
end = 6   # end is exclusive, matches ONNX Slice

def get_inputs():
    # Input shape: (N, C, H, W)
    x = torch.rand(batch_size, *input_shape)
    return [x]

def get_init_inputs():
    # Model __init__(axis, start, end)
    return [axis, start, end]
# Tensor / Tile
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a Tile operation on the input tensor.

    Given an input tensor x and a list of multiples, this model
    repeats x along each dimension according to `multiples`,
    equivalent to ONNX::Tile and torch.Tensor.repeat.

    Input:
        x: Tensor of shape (batch_size, *input_shape)

    Parameters:
        multiples (tuple or list of int):
            The repetition factor for each dimension. Its length must
            equal x.dim() (i.e. number of dimensions of the input).
    """

    def __init__(self, multiples):
        super(Model, self).__init__()
        self.multiples = tuple(int(m) for m in multiples)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == len(self.multiples), \
            f"Expected multiples length {len(self.multiples)} to match x.dim()={x.dim()}"
        return x.repeat(*self.multiples)


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

batch_size = 16
input_shape = (64, 64)   # C, W

multiples = (1, 2, 3)

def get_inputs():
    return [torch.rand(batch_size, *input_shape)]

def get_init_inputs():
    return [multiples]
# Tensor / TopK
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a Top-K operation, selecting the k largest elements.
    """
    def __init__(self, k: int = 10):
        super(Model, self).__init__()
        self.k = k
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Top-K operation to the input tensor along the last dimension.

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

        Returns:
            torch.Tensor: Top-K values tensor of shape (batch_size, k).
        """
        values, indices = torch.topk(x, self.k, dim=-1)
        return values

batch_size = 64
dim = 1024
K = 10

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

def get_init_inputs():
    return [K]  # k value for Top-K operation
# Tensor / TopK_3d
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs Top-K operation on 3D tensors.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Selects the k largest elements from 3D tensor along specified dimension.
        Useful for sequence or batched operations.
    """
    
    def __init__(self, k: int = 10, dim: int = -1):
        super(Model, self).__init__()
        self.k = k
        self.dim = dim
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, seq_len, dim]
        
        Returns:
            Top-K values tensor with shape [batch_size, seq_len, k]
        """
        values, _ = torch.topk(x, self.k, dim=self.dim)
        return values


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

batch_size = 16
seq_len = 64
dim = 128
k = 10
topk_dim = -1

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

def get_init_inputs():
    return [k, topk_dim]
# Tensor / TopK_dim0
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs Top-K operation along dimension 0.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Selects the k largest elements along dimension 0 (batch dimension).
        Useful for selecting top samples from a batch.
    """
    
    def __init__(self, k: int = 10, dim: int = 0):
        super(Model, self).__init__()
        self.k = k
        self.dim = dim
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, dim]
        
        Returns:
            Top-K values tensor with shape [k, dim]
        """
        values, _ = torch.topk(x, self.k, dim=self.dim)
        return values


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

batch_size = 32
dim = 256
k = 10
topk_dim = 0

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

def get_init_inputs():
    return [k, topk_dim]
# Tensor / TopK_indices
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs Top-K operation and returns only indices.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Selects the k largest elements along the specified dimension.
        Returns only the indices, not the values.
    """
    
    def __init__(self, k: int = 10, dim: int = -1):
        super(Model, self).__init__()
        self.k = k
        self.dim = dim
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, dim]
        
        Returns:
            Top-K indices tensor with shape [batch_size, k]
        """
        _, indices = torch.topk(x, self.k, dim=self.dim)
        return indices


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

batch_size = 32
dim = 256
k = 10
topk_dim = -1

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

def get_init_inputs():
    return [k, topk_dim]
# Tensor / TopK_sorted
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs Top-K operation with sorted output.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Selects the k largest elements and ensures they are sorted.
        Returns sorted values in descending order.
    """
    
    def __init__(self, k: int = 10, dim: int = -1, sorted: bool = True):
        super(Model, self).__init__()
        self.k = k
        self.dim = dim
        self.sorted = sorted
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [batch_size, dim]
        
        Returns:
            Sorted Top-K values tensor with shape [batch_size, k]
        """
        values, _ = torch.topk(x, self.k, dim=self.dim, sorted=self.sorted)
        return values


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

batch_size = 32
dim = 256
k = 10
topk_dim = -1
sorted_output = True

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

def get_init_inputs():
    return [k, topk_dim, sorted_output]
# Tensor / Trilu_lower
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that returns the lower triangular part of a matrix.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Returns the lower triangular part of the input matrix.
        Elements on and below the diagonal (offset=0) are kept.
        Elements above the diagonal are set to zero.
    """
    
    def __init__(self, upper: int = 0):
        """
        Initialize the Trilu model.
        
        Args:
            upper: 1 for upper triangle, 0 for lower triangle
        """
        super(Model, self).__init__()
        self.upper = upper
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [rows, cols]
        
        Returns:
            Output tensor with triangular part extracted
        """
        diagonal = 0
        if self.upper:
            return torch.triu(x, diagonal=diagonal)
        else:
            return torch.tril(x, diagonal=diagonal)


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

rows = 32
cols = 64
upper = 0

def get_inputs():
    x = torch.randint(0, 10, (rows, cols), dtype=torch.int64)
    return [x]

def get_init_inputs():
    return [upper]
# Tensor / Trilu_uppper
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that returns the lower triangular part of a matrix.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Returns the lower triangular part of the input matrix.
        Elements on and below the diagonal (offset=0) are kept.
        Elements above the diagonal are set to zero.
    """
    
    def __init__(self, upper: int = 0):
        """
        Initialize the Trilu model.
        
        Args:
            upper: 1 for upper triangle, 0 for lower triangle
        """
        super(Model, self).__init__()
        self.upper = upper
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with shape [rows, cols]
        
        Returns:
            Output tensor with triangular part extracted
        """
        diagonal = 0
        if self.upper:
            return torch.triu(x, diagonal=diagonal)
        else:
            return torch.tril(x, diagonal=diagonal)


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

rows = 32
cols = 64
upper = 1

def get_inputs():
    x = torch.randint(0, 10, (rows, cols), dtype=torch.int64)
    return [x]

def get_init_inputs():
    return [upper]
# Tensor / Unique
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional, List, Tuple, Any

class Model(nn.Module):
    def __init__(self, axis: int = 0, sorted: int = 1):
        super(Model, self).__init__()
        self.axis = axis
        self.sorted = sorted
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        axis = self.axis if hasattr(self, 'axis') else None
        sorted_out = self.sorted if hasattr(self, 'sorted') else 1
        
        if axis is None:
            output, inverse_indices, counts = torch.unique(x.flatten(), 
                                                          return_inverse=True, 
                                                          return_counts=True, 
                                                          sorted=bool(sorted_out))
            return output, inverse_indices, counts
        else:
            output, inverse_indices = torch.unique(x, dim=axis, return_inverse=True, sorted=bool(sorted_out))
            return output, inverse_indices

batch_size = 32
dim1 = 512
dim2 = 512

# ======== Example input configuration ========
def get_inputs():
    x = torch.randn(batch_size, dim1, dim2)
    return [x]

def get_init_inputs():
    return [0, 1]
# Tensor / Unique_sorted_with_axis
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Optional, Tuple, Union


class Model(nn.Module):
    
    def __init__(self, dim: Optional[int] = 0, sorted: bool = True):
        super(Model, self).__init__()
        self.dim = dim  
        self.sorted = sorted  
    
    def forward(self, X) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Forward pass
        
        Args:
            X: Tensor
        
        Returns:
            Tuple of (unique values, inverse indices, counts)
        """
        return torch.unique(
            X, 
            sorted=self.sorted,
            return_inverse=True, 
            return_counts=True,
            dim=self.dim
        )

batch_size = 64
dim = 512

# ======== Example input configuration ========
def get_inputs() -> List[torch.Tensor]:
    """
    Generate example inputs based on real ONNX test data
    
    Returns:
        List of input tensors
    """
    return [
        torch.randn(batch_size, dim),  # X
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return [0, True]  # [dim, sorted]