Jhnony/Mobilekernelbench huggingface dataset:
# Tensor / Unique_sorted_with_axis_3d
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] = 1, 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 = 16
dim1 = 64
dim2 = 128

# ======== 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, dim1, dim2),  # X
    ]


def get_init_inputs() -> List:
    """Get initialization inputs (if any)"""
    return [1, True]  # [dim, sorted]
# Tensor / Where
import torch
import torch.nn as nn
from typing import List


class Model(nn.Module):
    """
    A model that performs the Where (Select) operation.
    
    Number of inputs: 3
    Implementation type: direct
    
    Semantics:
        Selects elements from x or y based on condition.
        When condition is True, yield x, otherwise yield y.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, condition: torch.Tensor, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            condition: Boolean tensor indicating which elements to select
            x: First input tensor (selected when condition is True)
            y: Second input tensor (selected when condition is False)
        
        Returns:
            Output tensor with elements selected from x or y
        """
        return torch.where(condition, x, y)


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

batch_size = 32
features = 512
dim1 = 512

def get_inputs():
    condition = torch.rand(batch_size, features, dim1) > 0.5
    x = torch.randn(batch_size, features, dim1)
    y = torch.randn(batch_size, features, dim1)
    return [condition, x, y]

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


class Model(nn.Module):
    """
    A model that performs the Where operation with broadcasting.
    
    Number of inputs: 3
    Implementation type: direct
    
    Semantics:
        Selects elements from x or y based on condition.
        Demonstrates broadcasting where condition and inputs have different shapes.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, condition: torch.Tensor, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            condition: Boolean tensor with shape [batch_size, 1, 1]
            x: First input tensor with shape [batch_size, channels, features]
            y: Second input tensor with shape [batch_size, channels, features]
        
        Returns:
            Output tensor with elements selected from x or y
        """
        return torch.where(condition, x, y)


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

batch_size = 8
channels = 64
features = 512

def get_inputs():
    # Condition broadcasts across channels and features
    condition = torch.rand(batch_size, 1, 1) > 0.5
    x = torch.randn(batch_size, channels, features)
    y = torch.randn(batch_size, channels, features)
    return [condition, x, y]

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


class Model(nn.Module):
    """
    A model that performs the Where operation with scalar values.
    
    Number of inputs: 3
    Implementation type: direct
    
    Semantics:
        Selects elements from x or y based on condition.
        Demonstrates usage with scalar values that broadcast to match condition shape.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, condition: torch.Tensor, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            condition: Boolean tensor
            x: Scalar tensor (broadcasts to condition shape)
            y: Scalar tensor (broadcasts to condition shape)
        
        Returns:
            Output tensor with elements selected from x or y
        """
        return torch.where(condition, x, y)


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

batch_size = 4
height = 128
width = 128

def get_inputs():
    condition = torch.rand(batch_size, height, width) > 0.5
    # Scalar values that will broadcast
    x = torch.tensor(1.0)
    y = torch.tensor(0.0)
    return [condition, x, y]

def get_init_inputs():
    return []
# Trigonometry / Acos
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise arccosine (Acos) of the input tensor.

    Semantics:
        y = acos(x)
    """

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

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


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

batch_size = 16
input_shape = (128, 128)

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

def get_init_inputs():
    return []
# Trigonometry / Acosh
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 inverse hyperbolic cosine function element-wise.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Calculates the hyperbolic arccosine (acosh) of the input tensor element-wise.
        Input values must be in the range [1, inf).
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with values >= 1.0
        
        Returns:
            Output tensor after applying acosh function
        """
        return torch.acosh(x)


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

dim0 = 8
dim1 = 64
dim2 = 64

def get_inputs():
    x = torch.rand(dim0, dim1, dim2) + 1.0  # Range: [1.0, 2.0)
    return [x]

def get_init_inputs():
    return []
# Trigonometry / Asin
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise arcsine (Asin) of the input tensor.

    Semantics:
        y = asin(x)
    """

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

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


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

batch_size = 16
input_shape = (512, 256)

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

def get_init_inputs():
    return []
# Trigonometry / Asinh
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 inverse hyperbolic sine function element-wise.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Calculates the hyperbolic arcsine (asinh) of the input tensor element-wise.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying asinh function
        """
        return torch.asinh(x)


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

dim0 = 8
dim1 = 128
dim2 = 256

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

def get_init_inputs():
    return []
# Trigonometry / ATan
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that applies element-wise arctangent (ATan) to the input tensor.

    Semantics:
        y = atan(x)
    """

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

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


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

batch_size = 16
input_shape = (128, 512)  # you can adjust as needed

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

def get_init_inputs():
    return []
# Trigonometry / Atanh
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 inverse hyperbolic tangent function element-wise.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Calculates the hyperbolic arctangent (atanh) of the input tensor element-wise.
        Input values must be in the range (-1, 1).
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor with values in the range (-1, 1)
        
        Returns:
            Output tensor after applying atanh function
        """
        return torch.atanh(x)


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

dim0 = 8
dim1 = 64
dim2 = 64

def get_inputs():
    x = torch.rand(dim0, dim1, dim2) * 1.8 - 0.9  # Range: (-0.9, 0.9)
    return [x]

def get_init_inputs():
    return []
# Trigonometry / Cos
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise cosine (Cos) of the input tensor.

    Semantics:
        y = cos(x)
    """

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

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


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

batch_size = 16
input_shape = (256, 256)

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

def get_init_inputs():
    return []
# Trigonometry / Cosh
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 hyperbolic cosine function element-wise.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Calculates the hyperbolic cosine (cosh) of the input tensor element-wise.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying cosh function
        """
        return torch.cosh(x)


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

dim0 = 8
dim1 = 128
dim2 = 256

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

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


class Model(nn.Module):
    """
    A model that computes element-wise sine function.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes the sine of each element in the input tensor.
        Output has the same shape as input.
        Input values are interpreted as radians.
        
        output = sin(input)
    """
    
    def __init__(self):
        """
        Initialize the Sin model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            input: Input tensor (values in radians)
        
        Returns:
            Tensor with element-wise sine, values in range [-1, 1]
        """
        return torch.sin(input)


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

dim1 = 16
dim2 = 32
dim3 = 32

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

def get_init_inputs():
    return []
# Trigonometry / Sinh
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 hyperbolic sine function element-wise.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Calculates the hyperbolic sine (sinh) of the input tensor element-wise.
    """
    
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            x: Input tensor
        
        Returns:
            Output tensor after applying sinh function
        """
        return torch.sinh(x)


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

dim0 = 8
dim1 = 512
dim2 = 512

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

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


class Model(nn.Module):
    """
    A model that computes element-wise tangent function.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes the tangent of each element in the input tensor.
        Output has the same shape as input.
        Input values are interpreted as radians.
        
        output = tan(input) = sin(input) / cos(input)
        
    Note:
        Tangent has singularities at x = (2k+1)π/2 where k is an integer.
        At these points, tan(x) approaches ±∞.
    """
    
    def __init__(self):
        """
        Initialize the Tan model.
        
        Note: This operator has no learnable parameters.
        """
        super(Model, self).__init__()
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            input: Input tensor (values in radians)
        
        Returns:
            Tensor with element-wise tangent, values in range (-∞, +∞)
        """
        return torch.tan(input)


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

dim1 = 16
dim2 = 256
dim3 = 256

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

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


class Model(nn.Module):
    """
    A model that computes element-wise hyperbolic tangent function.
    
    Number of inputs: 1
    Implementation type: direct
    
    Semantics:
        Computes the hyperbolic tangent of each element in the input tensor.
        Output has the same shape as input.
        Output is bounded: tanh(x) ∈ (-1, 1)
        
        tanh(x) = (e^x - e^(-x)) / (e^x + e^(-x))
                = (e^(2x) - 1) / (e^(2x) + 1)
        
    Note:
        Unlike tan(x), tanh(x) has no singularities and is smooth everywhere.
        Commonly used as an activation function in neural networks.
    """
    
    def __init__(self):
        """
        Initialize the Tanh 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 shape
        
        Returns:
            Tensor with element-wise hyperbolic tangent, values in range (-1, 1)
        """
        return torch.tanh(x)


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

dim1 = 8
dim2 = 128
dim3 = 256

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

def get_init_inputs():
    return []
# Unary / Abs
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that applies element-wise absolute value (Abs) to the input tensor.

    Semantics:
        y = |x|
    """

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

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


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

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

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

def get_init_inputs():
    return []
# Unary / Ceil
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ceil(x)

batch_size = 8
dim1 = 128
dim2 = 256

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

def get_init_inputs():
    return []
# Unary / Erf
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Element-wise error function (Erf).

    Semantics:
        y = erf(x)
    """

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

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


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

batch_size = 16
input_shape = (64, 64)

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

def get_init_inputs():
    return []
# Unary / Exp
import torch
import torch.nn as nn

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

    Semantics:
        y = exp(x)
    """

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

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


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

batch_size = 16
input_shape = (512, 512)

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

def get_init_inputs():
    return []
# Unary / Floor
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.floor(x)

batch_size = 8
dim1 = 256
dim2 = 512

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

def get_init_inputs():
    return []
# Unary / Log
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):
        super(Model, self).__init__()
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return torch.log(input)

batch_size = 8
dim1 = 64
dim2 = 256

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

def get_init_inputs():
    return []
# Unary / Neg
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.neg(x)

batch_size = 8
dim1 = 256
dim2 = 256

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

def get_init_inputs():
    return []
# Unary / Reciprocal
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.reciprocal(x)

batch_size = 8
dim1 = 256
dim2 = 256

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

def get_init_inputs():
    return []
# Unary / Round
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.round(x)

batch_size = 8
dim1 = 512
dim2 = 512

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

def get_init_inputs():
    return []
# Unary / Sign
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):
        super(Model, self).__init__()
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return torch.sign(input)

batch_size = 8
dim1 = 128
dim2 = 256

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

def get_init_inputs():
    return []
# Unary / Sqrt
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):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.sqrt(x)

batch_size = 8
dim1 = 64
dim2 = 64

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

def get_init_inputs():
    return []