-- | Stateful token mixers from the pedagogic GPT-2 → linear → delta ladder
-- (waterloo_intern worklog / Kimi lineage).
--
-- Already in this package: full-sequence softmax MHA ('Circuit.LLM.GPT',
-- 'Circuit.LLM.Attention').  This module adds the /memory/ progression:
--
-- * softmax decode with an explicit KV cache
-- * linear attention with fixed-size state @S@ (and normalizer @z@)
-- * delta-rule write (overwrite what the key already stores)
-- * gated decay of @S@ (Mamba-style scalar, or per-channel vector for KDA-lite)
-- * chunked additive linear scan (matches sequential; training-shaped)
--
-- Shapes (single head, teaching form):
--
-- @
-- q, k, v  :: Matrix  -- T×d  (rows = time, cols = head dim)
-- S        :: Matrix  -- d×d  associative memory
-- @
--
-- Multi-head is a 'map' over heads later; keep one head honest first.
--
-- Pedagogic imprint: same blocks as the article, hmatrix BLAS, pure scans —
-- the discrete Process / decaying Stats story, but for attention state.
module Circuit.LLM.Mixer
  ( -- * Softmax + KV cache
    KvCache (..),
    emptyKv,
    softmaxAttnPrefill,
    softmaxAttnStep,

    -- * Linear attention (fixed state)
    LinearState (..),
    emptyLinear,
    eluPlus1,
    linearAttnPrefill,
    linearAttnStep,

    -- * Delta rule
    deltaAttnStep,
    deltaAttnPrefill,

    -- * Gated memory
    gatedDeltaStep,
    gatedDeltaPrefill,
    kdaLiteStep,

    -- * Chunked additive linear (matches sequential)
    chunkedLinearAttn,

    -- * Scans
    scanMixer,
  )
where

import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cmap,
    cols,
    fromLists,
    fromRows,
    konst,
    rows,
    scale,
    sumElements,
    toList,
    toRows,
    tr,
    (<>),
  )
import Numeric.LinearAlgebra qualified as LA
import Prelude hiding ((<>))

----------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------

-- | Causal upper-triangle mask: 0 on/below diagonal, -∞ above.
causalMask :: Int -> Matrix Double
causalMask :: Int -> Matrix Double
causalMask Int
n =
  [[Double]] -> Matrix Double
forall t. Element t => [[t]] -> Matrix t
fromLists
    [ [if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i then Double
0 else -(Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]

-- | Row-wise stable softmax.
softmaxRows :: Matrix Double -> Matrix Double
softmaxRows :: Matrix Double -> Matrix Double
softmaxRows Matrix Double
x = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [[Double] -> Vector Double
forall {a}. (Ord a, Floating a, Storable a) => [a] -> Vector a
soft (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
r) | Vector Double
r <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x]
  where
    soft :: [a] -> Vector a
soft [a]
xs =
      let m :: a
m = [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [a]
xs
          e :: [a]
e = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (\a
v -> a -> a
forall a. Floating a => a -> a
exp (a
v a -> a -> a
forall a. Num a => a -> a -> a
- a
m)) [a]
xs
          s :: a
s = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
e
       in [a] -> Vector a
forall a. Storable a => [a] -> Vector a
LA.fromList ((a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
s) [a]
e)

-- | Scaled dot-product attention (full sequence), causal.
softmaxAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
softmaxAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
softmaxAttnPrefill Matrix Double
q Matrix Double
k Matrix Double
v =
  let t :: Int
t = Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
q
      dk :: Double
dk = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
k) :: Double
      scores :: Matrix Double
scores = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
dk) (Matrix Double
q Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
k) Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Int -> Matrix Double
causalMask Int
t
      attn :: Matrix Double
attn = Matrix Double -> Matrix Double
softmaxRows Matrix Double
scores
   in Matrix Double
attn Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
v

----------------------------------------------------------------------
-- Softmax + KV cache
----------------------------------------------------------------------

-- | Growing key/value cache for decoder-style softmax attention.
data KvCache = KvCache
  { KvCache -> Matrix Double
kvK :: Matrix Double, -- T_past × d
    KvCache -> Matrix Double
kvV :: Matrix Double
  }
  deriving (KvCache -> KvCache -> Bool
(KvCache -> KvCache -> Bool)
-> (KvCache -> KvCache -> Bool) -> Eq KvCache
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: KvCache -> KvCache -> Bool
== :: KvCache -> KvCache -> Bool
$c/= :: KvCache -> KvCache -> Bool
/= :: KvCache -> KvCache -> Bool
Eq, Int -> KvCache -> ShowS
[KvCache] -> ShowS
KvCache -> String
(Int -> KvCache -> ShowS)
-> (KvCache -> String) -> ([KvCache] -> ShowS) -> Show KvCache
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KvCache -> ShowS
showsPrec :: Int -> KvCache -> ShowS
$cshow :: KvCache -> String
show :: KvCache -> String
$cshowList :: [KvCache] -> ShowS
showList :: [KvCache] -> ShowS
Show)

-- | Empty cache (0 rows).
emptyKv :: Int -> KvCache
emptyKv :: Int -> KvCache
emptyKv Int
d = Matrix Double -> Matrix Double -> KvCache
KvCache (Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
0, Int
d)) (Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
0, Int
d))

-- | One decode step: @q,k,v@ are 1×d (current token). Returns output 1×d and
-- extended cache.
softmaxAttnStep ::
  Matrix Double -> -- q 1×d
  Matrix Double -> -- k 1×d
  Matrix Double -> -- v 1×d
  KvCache ->
  (Matrix Double, KvCache)
softmaxAttnStep :: Matrix Double
-> Matrix Double
-> Matrix Double
-> KvCache
-> (Matrix Double, KvCache)
softmaxAttnStep Matrix Double
q Matrix Double
k Matrix Double
v KvCache
cache =
  let kAll :: Matrix Double
kAll = if Matrix Double -> Int
forall t. Matrix t -> Int
rows (KvCache -> Matrix Double
kvK KvCache
cache) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Matrix Double
k else KvCache -> Matrix Double
kvK KvCache
cache Matrix Double -> Matrix Double -> Matrix Double
forall t. Element t => Matrix t -> Matrix t -> Matrix t
LA.=== Matrix Double
k
      vAll :: Matrix Double
vAll = if Matrix Double -> Int
forall t. Matrix t -> Int
rows (KvCache -> Matrix Double
kvV KvCache
cache) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Matrix Double
v else KvCache -> Matrix Double
kvV KvCache
cache Matrix Double -> Matrix Double -> Matrix Double
forall t. Element t => Matrix t -> Matrix t -> Matrix t
LA.=== Matrix Double
v
      dk :: Double
dk = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
k) :: Double
      scores :: Matrix Double
scores = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
dk) (Matrix Double
q Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
kAll) -- 1 × T
      attn :: Matrix Double
attn = Matrix Double -> Matrix Double
softmaxRows Matrix Double
scores
      o :: Matrix Double
o = Matrix Double
attn Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
vAll
   in (Matrix Double
o, Matrix Double -> Matrix Double -> KvCache
KvCache Matrix Double
kAll Matrix Double
vAll)

----------------------------------------------------------------------
-- Linear attention (ELU+1 feature map, fixed S and z)
----------------------------------------------------------------------

-- | Fixed-size linear-attention memory: @S@ accumulates @φ(k)ᵀ v@ outer
-- products; @z@ accumulates @φ(k)@ for the normalizer.
data LinearState = LinearState
  { LinearState -> Matrix Double
linS :: Matrix Double, -- d×d
    LinearState -> Vector Double
linZ :: Vector Double -- d
  }
  deriving (LinearState -> LinearState -> Bool
(LinearState -> LinearState -> Bool)
-> (LinearState -> LinearState -> Bool) -> Eq LinearState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LinearState -> LinearState -> Bool
== :: LinearState -> LinearState -> Bool
$c/= :: LinearState -> LinearState -> Bool
/= :: LinearState -> LinearState -> Bool
Eq, Int -> LinearState -> ShowS
[LinearState] -> ShowS
LinearState -> String
(Int -> LinearState -> ShowS)
-> (LinearState -> String)
-> ([LinearState] -> ShowS)
-> Show LinearState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LinearState -> ShowS
showsPrec :: Int -> LinearState -> ShowS
$cshow :: LinearState -> String
show :: LinearState -> String
$cshowList :: [LinearState] -> ShowS
showList :: [LinearState] -> ShowS
Show)

emptyLinear :: Int -> LinearState
emptyLinear :: Int -> LinearState
emptyLinear Int
d = Matrix Double -> Vector Double -> LinearState
LinearState (Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
d, Int
d)) (Double -> Int -> Vector Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 Int
d)

-- | Feature map φ(x) = ELU(x)+1 (article / Katharopoulos linear attention).
eluPlus1 :: Matrix Double -> Matrix Double
eluPlus1 :: Matrix Double -> Matrix Double
eluPlus1 = (Double -> Double) -> Matrix Double -> Matrix Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
x -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 else Double -> Double
forall a. Floating a => a -> a
exp Double
x)

-- | One token of linear attention. @q,k,v@ are 1×d.
--
-- @
-- S ← S + φ(k)ᵀ v
-- z ← z + φ(k)
-- o  = (φ(q) S) / (φ(q)·z)
-- @
linearAttnStep ::
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  LinearState ->
  (Matrix Double, LinearState)
linearAttnStep :: Matrix Double
-> Matrix Double
-> Matrix Double
-> LinearState
-> (Matrix Double, LinearState)
linearAttnStep Matrix Double
q Matrix Double
k Matrix Double
v LinearState
st =
  let qf :: Matrix Double
qf = Matrix Double -> Matrix Double
eluPlus1 Matrix Double
q -- 1×d
      kf :: Matrix Double
kf = Matrix Double -> Matrix Double
eluPlus1 Matrix Double
k
      -- S is d×d storing sum_i φ(k_i)^T v_i  (outer: d×1 * 1×d)
      kCol :: Matrix Double
kCol = Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
kf -- d×1
      vRow :: Matrix Double
vRow = Matrix Double
v -- 1×d
      s' :: Matrix Double
s' = LinearState -> Matrix Double
linS LinearState
st Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ (Matrix Double
kCol Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
vRow)
      z' :: Vector Double
z' = LinearState -> Vector Double
linZ LinearState
st Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
kf
      -- o = φ(q) @ S  → 1×d; denom = φ(q)·z
      num :: Matrix Double
num = Matrix Double
qf Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
s' -- 1×d
      denom :: Double
denom = Matrix Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Matrix Double
qf Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
* Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
z') -- scalar as sum of elementwise
      o :: Matrix Double
o =
        if Double -> Double
forall a. Num a => a -> a
abs Double
denom Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12
          then Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
1, Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
q)
          else Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
denom) Matrix Double
num
   in (Matrix Double
o, Matrix Double -> Vector Double -> LinearState
LinearState Matrix Double
s' Vector Double
z')

-- | Prefill by scanning tokens (rows of q,k,v).
linearAttnPrefill ::
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, LinearState)
linearAttnPrefill :: Matrix Double
-> Matrix Double -> Matrix Double -> (Matrix Double, LinearState)
linearAttnPrefill Matrix Double
q Matrix Double
k Matrix Double
v =
  let d :: Int
d = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
q
      steps :: [(Vector Double, Vector Double, Vector Double)]
steps = [Vector Double]
-> [Vector Double]
-> [Vector Double]
-> [(Vector Double, Vector Double, Vector Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
q) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
k) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
v)
      ([Vector Double]
outs, LinearState
st) =
        (([Vector Double], LinearState)
 -> (Vector Double, Vector Double, Vector Double)
 -> ([Vector Double], LinearState))
-> ([Vector Double], LinearState)
-> [(Vector Double, Vector Double, Vector Double)]
-> ([Vector Double], LinearState)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \([Vector Double]
acc, LinearState
s) (Vector Double
qr, Vector Double
kr, Vector Double
vr) ->
              let (Matrix Double
o, LinearState
s') =
                    Matrix Double
-> Matrix Double
-> Matrix Double
-> LinearState
-> (Matrix Double, LinearState)
linearAttnStep
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
qr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
kr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
vr)
                      LinearState
s
               in ([Vector Double]
acc [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
++ [Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
o], LinearState
s')
          )
          ([], Int -> LinearState
emptyLinear Int
d)
          [(Vector Double, Vector Double, Vector Double)]
steps
   in ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
outs, LinearState
st)

----------------------------------------------------------------------
-- Delta rule
----------------------------------------------------------------------

-- | One delta-rule step (sequential form from the article).
--
-- @
-- v_old = k S          -- what this key already retrieves (k is 1×d, S is d×d:
--                      -- we use S as mapping d→d on row vectors: o = q <> S)
-- u     = β (v - v_old)
-- S'    = S + kᵀ u
-- o     = q S'
-- @
--
-- Convention: treat @S@ as right-acting on 1×d queries (@o = q <> S@), so
-- writes are @S += kᵀ <> u@ with @k,u@ as 1×d.
deltaAttnStep ::
  Double -> -- β write strength in (0,1] typically sigmoid
  Matrix Double -> -- q 1×d
  Matrix Double -> -- k 1×d
  Matrix Double -> -- v 1×d
  Matrix Double -> -- S d×d
  (Matrix Double, Matrix Double)
deltaAttnStep :: Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
deltaAttnStep Double
beta Matrix Double
q Matrix Double
k Matrix Double
v Matrix Double
s =
  let vOld :: Matrix Double
vOld = Matrix Double
k Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
s -- 1×d
      u :: Matrix Double
u = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
beta (Matrix Double
v Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
- Matrix Double
vOld)
      s' :: Matrix Double
s' = Matrix Double
s Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ (Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
k Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
u) -- d×d
      o :: Matrix Double
o = Matrix Double
q Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
s'
   in (Matrix Double
o, Matrix Double
s')

-- | Prefill with constant β.
deltaAttnPrefill ::
  Double ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, Matrix Double)
deltaAttnPrefill :: Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
deltaAttnPrefill Double
beta Matrix Double
q Matrix Double
k Matrix Double
v =
  let d :: Int
d = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
q
      steps :: [(Vector Double, Vector Double, Vector Double)]
steps = [Vector Double]
-> [Vector Double]
-> [Vector Double]
-> [(Vector Double, Vector Double, Vector Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
q) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
k) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
v)
      ([Vector Double]
outs, Matrix Double
sF) =
        (([Vector Double], Matrix Double)
 -> (Vector Double, Vector Double, Vector Double)
 -> ([Vector Double], Matrix Double))
-> ([Vector Double], Matrix Double)
-> [(Vector Double, Vector Double, Vector Double)]
-> ([Vector Double], Matrix Double)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \([Vector Double]
acc, Matrix Double
s) (Vector Double
qr, Vector Double
kr, Vector Double
vr) ->
              let (Matrix Double
o, Matrix Double
s') =
                    Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
deltaAttnStep
                      Double
beta
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
qr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
kr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
vr)
                      Matrix Double
s
               in ([Vector Double]
acc [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
++ [Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
o], Matrix Double
s')
          )
          ([], Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
d, Int
d))
          [(Vector Double, Vector Double, Vector Double)]
steps
   in ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
outs, Matrix Double
sF)

----------------------------------------------------------------------
-- Gated delta (scalar α) and KDA-lite (per-channel α)
----------------------------------------------------------------------

-- | Gated delta step: decay previous state, then delta-write.
--
-- @
-- S ← α S
-- then deltaAttnStep β
-- @
--
-- @α = 1@ → pure delta; @α = 0@ → forget all then write.
gatedDeltaStep ::
  Double -> -- α ∈ [0,1]
  Double -> -- β
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, Matrix Double)
gatedDeltaStep :: Double
-> Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
gatedDeltaStep Double
alpha Double
beta Matrix Double
q Matrix Double
k Matrix Double
v Matrix Double
s =
  Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
deltaAttnStep Double
beta Matrix Double
q Matrix Double
k Matrix Double
v (Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
alpha Matrix Double
s)

gatedDeltaPrefill ::
  Double ->
  Double ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, Matrix Double)
gatedDeltaPrefill :: Double
-> Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
gatedDeltaPrefill Double
alpha Double
beta Matrix Double
q Matrix Double
k Matrix Double
v =
  let d :: Int
d = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
q
      steps :: [(Vector Double, Vector Double, Vector Double)]
steps = [Vector Double]
-> [Vector Double]
-> [Vector Double]
-> [(Vector Double, Vector Double, Vector Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
q) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
k) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
v)
      ([Vector Double]
outs, Matrix Double
sF) =
        (([Vector Double], Matrix Double)
 -> (Vector Double, Vector Double, Vector Double)
 -> ([Vector Double], Matrix Double))
-> ([Vector Double], Matrix Double)
-> [(Vector Double, Vector Double, Vector Double)]
-> ([Vector Double], Matrix Double)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \([Vector Double]
acc, Matrix Double
s) (Vector Double
qr, Vector Double
kr, Vector Double
vr) ->
              let (Matrix Double
o, Matrix Double
s') =
                    Double
-> Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
gatedDeltaStep
                      Double
alpha
                      Double
beta
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
qr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
kr)
                      (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
vr)
                      Matrix Double
s
               in ([Vector Double]
acc [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
++ [Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
o], Matrix Double
s')
          )
          ([], Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
d, Int
d))
          [(Vector Double, Vector Double, Vector Double)]
steps
   in ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
outs, Matrix Double
sF)

-- | KDA-lite: per-channel decay @α :: Vector d@ (elementwise scale of each
-- column of @S@ before the delta write). Full KDA has richer parameterisation;
-- this imprints the /fine-grained forget/ idea.
kdaLiteStep ::
  Vector Double -> -- α length d
  Double -> -- β
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, Matrix Double)
kdaLiteStep :: Vector Double
-> Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
kdaLiteStep Vector Double
alpha Double
beta Matrix Double
q Matrix Double
k Matrix Double
v Matrix Double
s =
  let a :: [Double]
a = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
alpha
      sDecayed :: Matrix Double
sDecayed =
        [[Double]] -> Matrix Double
forall t. Element t => [[t]] -> Matrix t
fromLists
          [ [ Double
sij Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
a [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j)
            | (Int
j, Double
sij) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
row
            ]
          | [Double]
row <- Matrix Double -> [[Double]]
forall t. Element t => Matrix t -> [[t]]
LA.toLists Matrix Double
s
          ]
   in Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double)
deltaAttnStep Double
beta Matrix Double
q Matrix Double
k Matrix Double
v Matrix Double
sDecayed

----------------------------------------------------------------------
-- Chunked additive linear (training-shaped; matches sequential)
----------------------------------------------------------------------

-- | Chunked additive linear attention (article intermediate form).
-- Within chunk: causal @q (kᵀ v)@-style via masked scores on φ features;
-- across chunks: fold into @S@ and read @q S@.
--
-- For feature map = identity and no normalizer this is didactic; we use
-- φ = ELU+1 and still accumulate @S@ as in sequential linear (without @z@
-- in the chunk path for simplicity of the equality oracle — compare
-- unnormalized numerators).
chunkedLinearAttn ::
  Int -> -- chunk size C
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double -- T×d outputs (unnormalized φ(q)S path)
chunkedLinearAttn :: Int
-> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
chunkedLinearAttn Int
cSize Matrix Double
q Matrix Double
k Matrix Double
v =
  let t :: Int
t = Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
q
      d :: Int
d = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
q
      qf :: Matrix Double
qf = Matrix Double -> Matrix Double
eluPlus1 Matrix Double
q
      kf :: Matrix Double
kf = Matrix Double -> Matrix Double
eluPlus1 Matrix Double
k
      nChunks :: Int
nChunks = (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
cSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
cSize
      go :: Int -> Matrix Double -> [Vector Double] -> [Vector Double]
go Int
i Matrix Double
s [Vector Double]
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nChunks = [Vector Double]
acc
        | Bool
otherwise =
            let lo :: Int
lo = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cSize
                hi :: Int
hi = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
t ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cSize)
                n :: Int
n = Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo
                qC :: Matrix Double
qC = Int -> Int -> Matrix Double -> Matrix Double
subRows Int
lo Int
n Matrix Double
qf
                kC :: Matrix Double
kC = Int -> Int -> Matrix Double -> Matrix Double
subRows Int
lo Int
n Matrix Double
kf
                vC :: Matrix Double
vC = Int -> Int -> Matrix Double -> Matrix Double
subRows Int
lo Int
n Matrix Double
v
                -- inter: qC <> S  (n×d)
                oPrev :: Matrix Double
oPrev = Matrix Double
qC Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
s
                -- intra: causal attn on features
                scores :: Matrix Double
scores = (Matrix Double
qC Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
kC) Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
* Int -> Matrix Double
lowerTriOnes Int
n
                oCurr :: Matrix Double
oCurr = Matrix Double
scores Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
vC
                o :: Matrix Double
o = Matrix Double
oPrev Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
oCurr
                s' :: Matrix Double
s' = Matrix Double
s Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ (Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
kC Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
vC)
             in Int -> Matrix Double -> [Vector Double] -> [Vector Double]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Matrix Double
s' ([Vector Double]
acc [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
++ Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
o)
   in [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows (Int -> Matrix Double -> [Vector Double] -> [Vector Double]
go Int
0 (Double -> (Int, Int) -> Matrix Double
forall e d (c :: * -> *). Konst e d c => e -> d -> c e
konst Double
0 (Int
d, Int
d)) [])

subRows :: Int -> Int -> Matrix Double -> Matrix Double
subRows :: Int -> Int -> Matrix Double -> Matrix Double
subRows Int
lo Int
n Matrix Double
m = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows (Int -> [Vector Double] -> [Vector Double]
forall a. Int -> [a] -> [a]
take Int
n (Int -> [Vector Double] -> [Vector Double]
forall a. Int -> [a] -> [a]
drop Int
lo (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
m)))

-- | Strict lower-triangular inclusive ones (causal 0/1 mask as multiply).
lowerTriOnes :: Int -> Matrix Double
lowerTriOnes :: Int -> Matrix Double
lowerTriOnes Int
n =
  [[Double]] -> Matrix Double
forall t. Element t => [[t]] -> Matrix t
fromLists
    [ [if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i then Double
1 else Double
0 | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]

----------------------------------------------------------------------
-- Generic scan
----------------------------------------------------------------------

-- | Scan a step @(q,k,v,state) → (o,state')@ over rows.
scanMixer ::
  (Matrix Double -> Matrix Double -> Matrix Double -> s -> (Matrix Double, s)) ->
  s ->
  Matrix Double ->
  Matrix Double ->
  Matrix Double ->
  (Matrix Double, s)
scanMixer :: forall s.
(Matrix Double
 -> Matrix Double -> Matrix Double -> s -> (Matrix Double, s))
-> s
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, s)
scanMixer Matrix Double
-> Matrix Double -> Matrix Double -> s -> (Matrix Double, s)
step s
s0 Matrix Double
q Matrix Double
k Matrix Double
v =
  let steps :: [(Vector Double, Vector Double, Vector Double)]
steps = [Vector Double]
-> [Vector Double]
-> [Vector Double]
-> [(Vector Double, Vector Double, Vector Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
q) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
k) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
v)
      ([Vector Double]
outs, s
sF) =
        (([Vector Double], s)
 -> (Vector Double, Vector Double, Vector Double)
 -> ([Vector Double], s))
-> ([Vector Double], s)
-> [(Vector Double, Vector Double, Vector Double)]
-> ([Vector Double], s)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \([Vector Double]
acc, s
s) (Vector Double
qr, Vector Double
kr, Vector Double
vr) ->
              let (Matrix Double
o, s
s') = Matrix Double
-> Matrix Double -> Matrix Double -> s -> (Matrix Double, s)
step (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
qr) (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
kr) (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
vr) s
s
               in ([Vector Double]
acc [Vector Double] -> [Vector Double] -> [Vector Double]
forall a. [a] -> [a] -> [a]
++ [Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
o], s
s')
          )
          ([], s
s0)
          [(Vector Double, Vector Double, Vector Double)]
steps
   in ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
outs, s
sF)