import torch
import torch.nn as nn


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

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

        # Region 3: |x| >= 2 (Output 0)
        y_sat = torch.zeros_like(x)

        # Region 2: 1 < |x| < 2 (Output 0.5 * (2 - |x|)^2)
        y_transition = 0.5 * (2.0 - x_abs).pow(2)

        # Region 1: |x| <= 1 (Output 1 - x^2 / 2)
        y_center = 1.0 - x.pow(2) / 2.0

        # Combine:
        y_out = torch.where(
            x_abs <= 1.0,
            y_center,
            torch.where(
                x_abs < 2.0,
                y_transition,
                y_sat
            )
        )

        return y_out


batch_size = 128
feature_dim = 512


def get_inputs():
    x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return []