1import torch
 2from torch import nn
 3
 4
 5class TwoChannelFractalBlock(nn.Module):
 6    """Ternary two-channel renormalization block.
 7
 8    Inputs have shape [batch, 3, width], representing three child
 9    representations with separate neutral (a) and defect (b) channels.
10    """
11
12    def __init__(self, width, hidden=None, eps=1e-6):
13        super().__init__()
14        hidden = hidden or 2 * width
15        self.eps = eps
16        self.x_raw = nn.Parameter(torch.tensor(-1.0))
17        self.lambda_raw = nn.Parameter(torch.tensor(0.0))
18        self.align_a = nn.Linear(width, width, bias=False)
19        self.align_b = nn.Linear(width, width, bias=False)
20        self.defect = nn.Sequential(
21            nn.Linear(6 * width, hidden),
22            nn.GELU(),
23            nn.Linear(hidden, width),
24            nn.Tanh(),
25        )
26        self.gate = nn.Sequential(
27            nn.Linear(6 * width, width),
28            nn.Sigmoid(),
29        )
30
31    def forward(self, a, b):
32        if a.ndim != 3 or b.shape != a.shape or a.shape[1] != 3:
33            raise ValueError("a and b must both have shape [batch, 3, width]")
34        aa, bb = self.align_a(a), self.align_b(b)
35        lam = torch.nn.functional.softplus(self.lambda_raw)
36        x = torch.nn.functional.softplus(self.x_raw)
37        neutral = lam**3 * aa[:, 0] * aa[:, 1] * aa[:, 2]
38        neutral = neutral + 2.0 * x**3 * bb[:, 0] * bb[:, 1] * bb[:, 2]
39        neutral = neutral / (neutral.square().mean(-1, keepdim=True).add(self.eps).sqrt())
40        flat = torch.cat((aa.reshape(a.shape[0], -1), bb.reshape(b.shape[0], -1)), dim=-1)
41        proposed = self.defect(flat)
42        gate = self.gate(flat)
43        defect = gate * proposed + (1.0 - gate) * bb.mean(dim=1)
44        return neutral, defect
45
46    def coefficients(self):
47        return (
48            torch.nn.functional.softplus(self.x_raw).item(),
49            torch.nn.functional.softplus(self.lambda_raw).item(),
50        )