import torch

def memory_forward(x_seq, Wq, Wk, Wv, W0, eta_w, eta_max=1.0):
    """One sequence: x_seq [T, d], T > 0.
    Wq/Wk [d, dk]; Wv [d, dv]; W0 [dk, dv]; eta_w [d].
    All tensors use a compatible floating dtype and device.
    """
    W = W0.clone()
    outputs = []
    for x in x_seq:
        q, k, v = x @ Wq, x @ Wk, x @ Wv
        eta = eta_max * torch.sigmoid(x @ eta_w)
        predicted = k @ W
        error = predicted - v
        grad_W = torch.outer(k, error)
        W = W - eta * grad_W
        outputs.append(q @ W)
    return torch.stack(outputs), W

# Outer training: send outputs through the rest of the LM;
# compute next-token cross-entropy, then loss.backward()
# and take an optimizer step on the slow parameters.
# Inference: torch.no_grad() is valid for THIS analytic
# implementation. The arithmetic memory write still runs.

import torch

def gated_memory_forward(x_seq, Wq, Wk, Wv, W0,
                         eta_w, mu_w, alpha_w, eta_max=1.0):
    """One sequence, same shapes as Linear GD.
    eta_w / mu_w / alpha_w: [d].
    Toy linear memory; not a full Titans layer.
    """
    W = W0.clone()
    U = torch.zeros_like(W)
    outputs = []
    for x in x_seq:
        q, k, v = x @ Wq, x @ Wk, x @ Wv
        eta = eta_max * torch.sigmoid(x @ eta_w)
        mu = torch.sigmoid(x @ mu_w)
        alpha = torch.sigmoid(x @ alpha_w)
        grad_W = torch.outer(k, k @ W - v)
        U = mu * U - eta * grad_W
        W = (1.0 - alpha) * W + U
        outputs.append(q @ W)
    return torch.stack(outputs), W, U

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

class TinyMemoryLM(nn.Module):
    """Teaching model: embedding → linear fast memory → LM head."""
    def __init__(self, vocab_size, d, dk, dv):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d)
        self.Wq = nn.Parameter(0.02 * torch.randn(d, dk))
        self.Wk = nn.Parameter(0.02 * torch.randn(d, dk))
        self.Wv = nn.Parameter(0.02 * torch.randn(d, dv))
        self.W0 = nn.Parameter(torch.zeros(dk, dv))
        self.eta_w = nn.Parameter(torch.zeros(d))
        self.eta_max = 1.0
        self.head = nn.Linear(dv, vocab_size)

def train_batches(model, dataloader, optimizer):
    """Each batch: long token IDs [B, T+1], B,T > 0.
    Equal-length independent sequences, no padding.
    Optimizer owns model.parameters(); batches share its device.
    """
    model.train()
    for n, batch in enumerate(dataloader):
        optimizer.zero_grad(set_to_none=True)
        losses = []
        for ids in batch:
            W = model.W0.clone()
            x_seq = model.embed(ids[:-1])
            for t, x in enumerate(x_seq):
                k, v = x @ model.Wk, x @ model.Wv
                q = x @ model.Wq
                eta = model.eta_max * torch.sigmoid(x @ model.eta_w)
                g = torch.outer(k, k @ W - v)
                W = W - eta * g
                logits = model.head(q @ W)
                losses.append(F.cross_entropy(
                    logits.unsqueeze(0), ids[t + 1].reshape(1)))
        loss = torch.stack(losses).mean()
        loss.backward()
        optimizer.step()

import torch

@torch.no_grad()
def read_prefix(model, ids):
    """ids: nonempty observed prefix [N], on model's device.
    Uses TinyMemoryLM and memory_forward from the other tabs.
    """
    model.eval()
    x_seq = model.embed(ids)
    outputs, W_final = memory_forward(
        x_seq, model.Wq, model.Wk, model.Wv,
        model.W0, model.eta_w, model.eta_max)
    next_logits = model.head(outputs[-1])
    return next_logits, W_final

# No next-token labels, loss.backward(), or optimizer.step().
# Fast W still changes at each prefix position.
# This helper replays the prefix from W0 on every call.