1import torch
 2from torch import nn
 3
 4
 5class FractionalMemory(nn.Module):
 6    """Causal positive exponential-mixture memory, accepting [T, B, D]."""
 7    def __init__(self, d_model: int, context_length: int, n_states: int = 8,
 8                 min_tau: float = 1.0, residual: bool = True):
 9        super().__init__()
10        if n_states < 1 or context_length < 1:
11            raise ValueError("n_states and context_length must be positive")
12        taus = torch.logspace(torch.log10(torch.tensor(float(min_tau))),
13                              torch.log10(torch.tensor(float(context_length))),
14                              n_states)
15        self.register_buffer("rho", torch.exp(-1.0 / taus))
16        self.logits = nn.Parameter(torch.zeros(n_states))
17        self.in_proj = nn.Linear(d_model, d_model)
18        self.out_proj = nn.Linear(d_model, d_model)
19        self.residual = residual
20
21    def forward(self, x: torch.Tensor) -> torch.Tensor:
22        if x.ndim != 3:
23            raise ValueError("expected [T, B, D]")
24        _, b, d = x.shape
25        u = self.in_proj(x)
26        state = torch.zeros(self.rho.numel(), b, d, device=x.device, dtype=x.dtype)
27        w = torch.softmax(self.logits, dim=0).to(dtype=x.dtype)
28        rho = self.rho.to(device=x.device, dtype=x.dtype).view(-1, 1, 1)
29        ys = []
30        for xt in u:
31            state = rho * state + (1.0 - rho) * xt.unsqueeze(0)
32            ys.append((w.view(-1, 1, 1) * state).sum(dim=0))
33        y = self.out_proj(torch.stack(ys, dim=0))
34        return x + y if self.residual else y