import torch
import torch.nn as nn

class OperatorLayer(nn.Module):
    def __init__(self, dim, op_set=['mul', 'max', 'softmax', 'entropic_fusion']):
        super().__init__()
        self.controller = nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, len(op_set)))
        self.W = nn.Parameter(torch.randn(dim, dim))
        self.op_set = op_set
        
    def forward(self, h):
        # h: (batch, dim)
        entropy = -torch.sum(torch.softmax(h.abs(), dim=0) * torch.log_softmax(h.abs(), dim=0), dim=0)
        logits = self.controller(entropy)  # (dim, n_ops)
        op_probs = torch.softmax(logits / 0.5, dim=-1)  # temperature=0.5
        
        # Apply weighted combination of operators (differentiable)
        out = 0
        for i, op_name in enumerate(self.op_set):
            op = get_operator(op_name)
            out = out + op_probs[:, i].unsqueeze(1) * op(self.W * h)
        return out