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


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.epsilon = 1e-5

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_flat = x.flatten()

        median = torch.quantile(x_flat, 0.5)

        q1 = torch.quantile(x_flat, 0.25)
        q3 = torch.quantile(x_flat, 0.75)
        iqr = q3 - q1

        z_robust = (x - median) / (iqr + self.epsilon)
        gate = torch.sigmoid(x)

        return z_robust * gate


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 []