import torch
import torch.nn as nn
import torch.nn.functional as F

N, C, H, W = 16, 16, 64, 64
LAMBDA = 0.5


class Softshrink(nn.Module):

    def __init__(self, lambd=LAMBDA):
        super().__init__()
        self.lambd = lambd

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        abs_val_minus_lambda = torch.abs(input) - self.lambd

        thresholded_magnitude = torch.relu(abs_val_minus_lambda)

        return torch.sign(input) * thresholded_magnitude


class Model(nn.Module):
    def __init__(self, lambd=LAMBDA):
        super().__init__()
        self.op = Softshrink(lambd=lambd)

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


# --- 辅助函数 ---

def get_inputs():
    torch.manual_seed(42)
    x = torch.randn(N, C, H, W, dtype=torch.float32) * 2.0
    return [x]


def get_init_inputs():
    return [LAMBDA]