43import math
44from typing import Optional
45
46import torch
47from torch import nn
48
49from labml_helpers.module import Module
50from labml_nn.transformers.feed_forward import FeedForward
51from labml_nn.transformers.mha import PrepareForMultiHeadAttention
52from labml_nn.utils import clone_module_list

55class FeedbackAttention(Module):

66    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1, *,
67                 is_kv_precomputed: bool = False):

75        super().__init__()

78        self.d_k = d_model // heads

80        self.heads = heads

83        self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)

85        if not is_kv_precomputed:
86            self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
87            self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=True)

89        else:
90            self.key = None
91            self.value = None

94        self.output = nn.Linear(d_model, d_model)

96        self.dropout = nn.Dropout(dropout_prob)

98        self.scale = 1 / math.sqrt(self.d_k)

101        self.softmax = nn.Softmax(dim=0)

104        self.P = 2 ** 12

107        self.key_pos_embeddings = nn.Parameter(torch.zeros((self.P, heads, self.d_k)), requires_grad=True)

109        self.key_pos_bias = nn.Parameter(torch.zeros((self.P, heads)), requires_grad=True)

111        self.query_pos_bias = nn.Parameter(torch.zeros((heads, self.d_k)), requires_grad=True)

114        self.attn = None

116    def get_scores(self, query: torch.Tensor, key: torch.Tensor):

144        key_pos_emb = self.key_pos_embeddings[-key.shape[0]:]

146        query_pos_bias = self.query_pos_bias[None, :, :]

148        key_pos_bias = self.key_pos_bias[-key.shape[0]:]

151        ac = torch.einsum('bhd,jbhd->jbh', query + query_pos_bias, key)

153        bd = torch.einsum('bhd,jhd->jbh', query, key_pos_emb) + key_pos_bias[:, None, :]

156        return ac + bd

158    def forward(self, *,
159                query: torch.Tensor,
160                key: torch.Tensor,
161                value: torch.Tensor):

170        query = self.query(query)
171        if self.key:
172            key = self.key(key)
173        if self.value:
174            value = self.value(value)

178        scores = self.get_scores(query, key)

181        scores *= self.scale

184        attn = self.softmax(scores)

187        attn = self.dropout(attn)

190        x = torch.einsum("jbh,jbhd->bhd", attn, value)

193        x = x.reshape(x.shape[0], -1)

196        return self.output(x)

199class FeedbackTransformerLayer(Module):

206    def __init__(self, *,
207                 d_model: int,
208                 attn: FeedbackAttention,
209                 feed_forward: FeedForward,
210                 dropout_prob: float):

217        super().__init__()

219        self.size = d_model

221        self.attn = attn
222        self.feed_forward = feed_forward
223        self.dropout = nn.Dropout(dropout_prob)

226        self.norm_self_attn = nn.LayerNorm([d_model])
227        self.norm_ff = nn.LayerNorm([d_model])

229    def forward(self, *,
230                x: torch.Tensor,
231                key: Optional[torch.Tensor],
232                value: Optional[torch.Tensor]):

234        if key is not None:

236            z = self.norm_self_attn(x)

238            self_attn = self.attn(query=z, key=key, value=value)

240            x = x + self.dropout(self_attn)

243        z = self.norm_ff(x)

245        ff = self.feed_forward(z)

247        x = x + self.dropout(ff)

250        return x

253class FeedbackTransformer(Module):

258    def __init__(self, layer: FeedbackTransformerLayer, n_layers: int):

264        super().__init__()

266        self.layers = clone_module_list(layer, n_layers)

268        self.norm = nn.LayerNorm([layer.size])

271        self.weights = nn.Parameter(torch.ones(n_layers + 1), requires_grad=True)

273        self.softmax = nn.Softmax(0)

275    def forward(self, x_seq: torch.Tensor):

281        x_seq = torch.unbind(x_seq, dim=0)

283        res = []

285        mem = []

287        for x in x_seq:

289            layer_outputs = [x]

292            mem_tensor = torch.stack(mem) if mem else None

295            for layer in self.layers:

297                x = layer(x=x, key=mem_tensor, value=mem_tensor)

299                layer_outputs.append(x)

302            layer_outputs = torch.stack(layer_outputs)

304            mem.append(torch.einsum('lbd,l->bd', layer_outputs, self.softmax(self.weights)))

306            res.append(x)

309        res = torch.stack(res)

311        return self.norm(res)

318class StackFunction(torch.autograd.Function):

330    @staticmethod
331    def forward(ctx, memory, memory_grad, last, n):

343        ctx._mem_grad = memory_grad

345        ctx._n = n

347        return memory[:n + 1]

349    @staticmethod
350    def backward(ctx, grad_output):

358        n = ctx._n

360        memory_grad = ctx._mem_grad

362        memory_grad[:n + 1] += grad_output

364        return None, None, memory_grad[n], None

367class Stack:

374    def __init__(self, max_len: int):

378        self.max_len = max_len
379        self.memory = None
380        self.memory_grad = None
381        self.last = None
382        self.n = -1
383        self.last_get_n = -1

385    def append(self, n: int, value: torch.Tensor):

393        assert n == 0 or self.last_get_n == n - 1, f"{n}, {self.last_get_n}"

396        with torch.no_grad():

398            if self.memory is None or self.memory.shape[1:] != value.shape:

400                assert n == 0

402                self.memory = value.new_zeros(self.max_len, *value.shape, requires_grad=False)

404                self.memory_grad = value.new_zeros(self.memory.shape, requires_grad=False)

409            elif n == 0:

411                self.memory_grad.fill_(0.)

414            self.memory.data[n] = value.detach()

416            self.n = n

421        self.last = value

423    def get(self):

430        self.last_get_n = self.n

433        return StackFunction.apply(self.memory, self.memory_grad, self.last, self.n)

435    def free(self):

440        self.memory = None
441        self.memory_grad = None
442        self.last = None

445class FeedbackTransformerKV(Module):

452    def __init__(self, layer: FeedbackTransformerLayer, n_layers: int, d_model: int, heads: int):

460        super().__init__()

462        self.layers = clone_module_list(layer, n_layers)

464        self.norm = nn.LayerNorm([layer.size])

467        self.weights = nn.Parameter(torch.ones(n_layers + 1), requires_grad=True)

469        self.softmax = nn.Softmax(0)

472        d_k = d_model // heads

474        self.key = PrepareForMultiHeadAttention(d_model, heads, d_k, bias=False)

476        self.value = PrepareForMultiHeadAttention(d_model, heads, d_k, bias=False)

479        self.mem_key = Stack(512)

481        self.mem_value = Stack(512)

483    def forward(self, x_seq: torch.Tensor):

489        x_seq = torch.unbind(x_seq, dim=0)

491        res = []

493        for step, x in enumerate(x_seq):

495            layer_outputs = [x]

498            key_tensor = None
499            value_tensor = None

501            if step > 0:
502                key_tensor = self.mem_key.get()
503                value_tensor = self.mem_value.get()

506            for layer in self.layers:

508                x = layer(x=x, key=key_tensor, value=value_tensor)

510                layer_outputs.append(x)

513            layer_outputs = torch.stack(layer_outputs)

515            mem = torch.einsum('lbd,l->bd', layer_outputs, self.softmax(self.weights))

517            self.mem_key.append(step, self.key(mem))

519            self.mem_value.append(step, self.value(mem))

521            res.append(x)

524        res = torch.stack(res)

526        return self.norm(res)

528    def free(self):
529        self.mem_key.free()
530        self.mem_value.free()