{-# LANGUAGE OverloadedStrings #-}

-- | Local compositional reverse-mode AD over hmatrix.
--
-- The idea is lifted from circuits-ad's 'Diff' arrow, but specialised to
-- dense 'hmatrix' matrices/vectors so that the heavy linear algebra still
-- calls BLAS.
--
-- Both causal (GPT-style) and bidirectional (BERT-style) transformer bodies
-- are provided.  They share the same block structure and differ only in the
-- attention mask that is added to the raw QK^T scores.
module Circuit.LLM.Diff
  ( -- * Differentiable operations
    DiffP,
    primForward,
    primBackward,
    (@.),
    residual,
    splitP,
    joinP,

    -- * Primitive layers
    linearP,
    geluP,
    softmaxP,
    layerNormP,
    multiHeadAttentionP,

    -- * Transformer block
    BlockParams (..),
    blockParamsFromBlock,
    blockDiffP,
    gptBlockDiffP,
    bertBlockDiffP,

    -- * Full model bodies
    GptParams (..),
    gptParamsFromModel,
    gptDiffP,
    bertDiffP,

    -- * Forward helpers (also used by Backprop)
    geluF,
    softmaxStableF,
    layerNormF,
    concatColsF,
    broadcastBias,
    causalMaskF,
    accumulateCols,
    subMatrixW,
    zeroMatrix,
  )
where

import Circuit.Diff.Param (TensorPrim (..))
import Circuit.Diff.Param qualified as ADP
import Circuit.LLM.GPT (FeedForward (..), Gpt (..), GptConfig (..), TransformerBlock (..))
import Data.List (foldl1')
import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cmap,
    cols,
    fromList,
    fromLists,
    fromRows,
    maxElement,
    reshape,
    rows,
    scale,
    sumElements,
    toList,
    toRows,
    tr,
    (|||),
  )
import Numeric.LinearAlgebra qualified as LA

----------------------------------------------------------------------
-- Compositional differentiable arrow
----------------------------------------------------------------------

-- | A differentiable operation with parameters @p@, input @a@ and output @b@.
--
-- Running forward produces the output.  The backward pass, given the output
-- cotangent, produces the input cotangent and parameter gradients.
--
-- This is exactly the same shape as 'Circuit.Diff.Param.TensorPrim' from
-- circuits-ad; we use a local type synonym so the rest of the module keeps
-- its paired-parameter composition style.
type DiffP = TensorPrim

-- | Sequential composition: @f '@.' g@ means "first @g@, then @f@".
infixr 9 @.

(@.) :: DiffP p2 b c -> DiffP p1 a b -> DiffP (p1, p2) a c
DiffP p2 b c
f2 @. :: forall p2 b c p1 a.
DiffP p2 b c -> DiffP p1 a b -> DiffP (p1, p2) a c
@. DiffP p1 a b
f1 =
  TensorPrim
    { primForward :: (p1, p2) -> a -> c
primForward = \(p1
p1, p2
p2) a
a -> DiffP p2 b c -> p2 -> b -> c
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p2 b c
f2 p2
p2 (DiffP p1 a b -> p1 -> a -> b
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p1 a b
f1 p1
p1 a
a),
      primBackward :: (p1, p2) -> a -> c -> (a, (p1, p2))
primBackward = \(p1
p1, p2
p2) a
a c
dc ->
        let b :: b
b = DiffP p1 a b -> p1 -> a -> b
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p1 a b
f1 p1
p1 a
a
            (b
db, p2
dp2) = DiffP p2 b c -> p2 -> b -> c -> (b, p2)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p2 b c
f2 p2
p2 b
b c
dc
            (a
da, p1
dp1) = DiffP p1 a b -> p1 -> a -> b -> (a, p1)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p1 a b
f1 p1
p1 a
a b
db
         in (a
da, (p1
dp1, p2
dp2))
    }

-- | Add a residual connection around an operation.
--   forward:  y = x + op(x)
--   backward: dx = dy + dOp
residual :: (Num a) => DiffP p a a -> DiffP p a a
residual :: forall a p. Num a => DiffP p a a -> DiffP p a a
residual DiffP p a a
op =
  TensorPrim
    { primForward :: p -> a -> a
primForward = \p
p a
a -> a
a a -> a -> a
forall a. Num a => a -> a -> a
+ DiffP p a a -> p -> a -> a
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p a a
op p
p a
a,
      primBackward :: p -> a -> a -> (a, p)
primBackward = \p
p a
a a
dy ->
        let (a
daOp, p
dp) = DiffP p a a -> p -> a -> a -> (a, p)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p a a
op p
p a
a a
dy
         in (a
dy a -> a -> a
forall a. Num a => a -> a -> a
+ a
daOp, p
dp)
    }

-- | Split a parameter tuple so the left and right halves can be used
--   by independent parallel branches.
splitP :: DiffP p1 a1 b1 -> DiffP p2 a2 b2 -> DiffP (p1, p2) (a1, a2) (b1, b2)
splitP :: forall p1 a1 b1 p2 a2 b2.
DiffP p1 a1 b1
-> DiffP p2 a2 b2 -> DiffP (p1, p2) (a1, a2) (b1, b2)
splitP DiffP p1 a1 b1
f1 DiffP p2 a2 b2
f2 =
  TensorPrim
    { primForward :: (p1, p2) -> (a1, a2) -> (b1, b2)
primForward = \(p1
p1, p2
p2) (a1
a1, a2
a2) -> (DiffP p1 a1 b1 -> p1 -> a1 -> b1
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p1 a1 b1
f1 p1
p1 a1
a1, DiffP p2 a2 b2 -> p2 -> a2 -> b2
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p2 a2 b2
f2 p2
p2 a2
a2),
      primBackward :: (p1, p2) -> (a1, a2) -> (b1, b2) -> ((a1, a2), (p1, p2))
primBackward = \(p1
p1, p2
p2) (a1
a1, a2
a2) (b1
db1, b2
db2) ->
        let (a1
da1, p1
dp1) = DiffP p1 a1 b1 -> p1 -> a1 -> b1 -> (a1, p1)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p1 a1 b1
f1 p1
p1 a1
a1 b1
db1
            (a2
da2, p2
dp2) = DiffP p2 a2 b2 -> p2 -> a2 -> b2 -> (a2, p2)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p2 a2 b2
f2 p2
p2 a2
a2 b2
db2
         in ((a1
da1, a2
da2), (p1
dp1, p2
dp2))
    }

-- | Pair two operations that share the same input type but produce
--   independent outputs.  This is useful for attention Q/K/V from one x.
joinP :: (Num a) => DiffP p1 a b1 -> DiffP p2 a b2 -> DiffP (p1, p2) a (b1, b2)
joinP :: forall a p1 b1 p2 b2.
Num a =>
DiffP p1 a b1 -> DiffP p2 a b2 -> DiffP (p1, p2) a (b1, b2)
joinP DiffP p1 a b1
f1 DiffP p2 a b2
f2 =
  TensorPrim
    { primForward :: (p1, p2) -> a -> (b1, b2)
primForward = \(p1
p1, p2
p2) a
a -> (DiffP p1 a b1 -> p1 -> a -> b1
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p1 a b1
f1 p1
p1 a
a, DiffP p2 a b2 -> p2 -> a -> b2
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP p2 a b2
f2 p2
p2 a
a),
      primBackward :: (p1, p2) -> a -> (b1, b2) -> (a, (p1, p2))
primBackward = \(p1
p1, p2
p2) a
a (b1
db1, b2
db2) ->
        let (a
da1, p1
dp1) = DiffP p1 a b1 -> p1 -> a -> b1 -> (a, p1)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p1 a b1
f1 p1
p1 a
a b1
db1
            (a
da2, p2
dp2) = DiffP p2 a b2 -> p2 -> a -> b2 -> (a, p2)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP p2 a b2
f2 p2
p2 a
a b2
db2
         in (a
da1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
da2, (p1
dp1, p2
dp2))
    }

----------------------------------------------------------------------
-- Primitive layers
----------------------------------------------------------------------

-- | Linear layer: y = xW + b
linearP :: DiffP (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP :: DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP =
  TensorPrim
    { primForward :: (Matrix Double, Vector Double) -> Matrix Double -> Matrix Double
primForward = \(Matrix Double
w, Vector Double
b) Matrix Double
x -> Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
w Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Vector Double -> Int -> Matrix Double
broadcastBias Vector Double
b (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
x),
      primBackward :: (Matrix Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Matrix Double, Vector Double))
primBackward = \(Matrix Double
w, Vector Double
_) Matrix Double
x Matrix Double
dy ->
        let dx :: Matrix Double
dx = Matrix Double
dy Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
w
            dw :: Matrix Double
dw = Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
dy
            db :: Vector Double
db = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
col | Vector Double
col <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
dy]
         in (Matrix Double
dx, (Matrix Double
dw, Vector Double
db))
    }

-- | GELU activation.  Migrated to use @circuits-ad:DiffP@.
geluP :: ADP.DiffP () (Matrix Double) (Matrix Double)
geluP :: DiffP () (Matrix Double) (Matrix Double)
geluP =
  TensorPrim () (Matrix Double) (Matrix Double)
-> DiffP () (Matrix Double) (Matrix Double)
forall p a b. TensorPrim p a b -> DiffP p a b
ADP.fromPrim (TensorPrim () (Matrix Double) (Matrix Double)
 -> DiffP () (Matrix Double) (Matrix Double))
-> TensorPrim () (Matrix Double) (Matrix Double)
-> DiffP () (Matrix Double) (Matrix Double)
forall a b. (a -> b) -> a -> b
$
    ADP.TensorPrim
      { primForward :: () -> Matrix Double -> Matrix Double
ADP.primForward = \()
_ Matrix Double
x -> Matrix Double -> Matrix Double
geluF Matrix Double
x,
        primBackward :: () -> Matrix Double -> Matrix Double -> (Matrix Double, ())
ADP.primBackward = \()
_ Matrix Double
x Matrix Double
dy -> (Matrix Double -> Matrix Double -> Matrix Double
geluBwd Matrix Double
x Matrix Double
dy, ())
      }

-- | Row-wise softmax.
softmaxP :: DiffP () (Matrix Double) (Matrix Double)
softmaxP :: TensorPrim () (Matrix Double) (Matrix Double)
softmaxP =
  TensorPrim
    { primForward :: () -> Matrix Double -> Matrix Double
primForward = \()
_ Matrix Double
x -> Matrix Double -> Matrix Double
softmaxStableF Matrix Double
x,
      primBackward :: () -> Matrix Double -> Matrix Double -> (Matrix Double, ())
primBackward = \()
_ Matrix Double
x Matrix Double
dy ->
        let probs :: Matrix Double
probs = Matrix Double -> Matrix Double
softmaxStableF Matrix Double
x
         in (Matrix Double -> Matrix Double -> Matrix Double
forall {t}. (Element t, Num t) => Matrix t -> Matrix t -> Matrix t
softmaxBwd Matrix Double
probs Matrix Double
dy, ())
    }
  where
    softmaxBwd :: Matrix t -> Matrix t -> Matrix t
softmaxBwd Matrix t
probs Matrix t
gradOut =
      [Vector t] -> Matrix t
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector t] -> Matrix t) -> [Vector t] -> Matrix t
forall a b. (a -> b) -> a -> b
$ (Vector t -> Vector t -> Vector t)
-> [Vector t] -> [Vector t] -> [Vector t]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Vector t -> Vector t -> Vector t
forall {a}. (Storable a, Num a) => Vector a -> Vector a -> Vector a
softmaxRowBwd (Matrix t -> [Vector t]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix t
probs) (Matrix t -> [Vector t]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix t
gradOut)
      where
        softmaxRowBwd :: Vector a -> Vector a -> Vector a
softmaxRowBwd Vector a
p Vector a
go =
          let pv :: [a]
pv = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
p; gov :: [a]
gov = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
go; dot :: a
dot = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
pv [a]
gov)
           in [a] -> Vector a
forall a. Storable a => [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
pi_ a
goi -> a
pi_ a -> a -> a
forall a. Num a => a -> a -> a
* (a
goi a -> a -> a
forall a. Num a => a -> a -> a
- a
dot)) [a]
pv [a]
gov

-- | Layer normalisation.
layerNormP :: Double -> DiffP (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
layerNormP :: Double
-> DiffP
     (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
layerNormP Double
eps =
  TensorPrim
    { primForward :: (Vector Double, Vector Double) -> Matrix Double -> Matrix Double
primForward = \(Vector Double
gamma, Vector Double
beta) Matrix Double
x -> Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
x Vector Double
gamma Vector Double
beta Double
eps,
      primBackward :: (Vector Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Vector Double, Vector Double))
primBackward = \(Vector Double
gamma, Vector Double
beta) Matrix Double
x Matrix Double
dy ->
        let (Matrix Double
dx, Vector Double
dgamma, Vector Double
dbeta) = Matrix Double
-> Vector Double
-> Vector Double
-> Double
-> Matrix Double
-> (Matrix Double, Vector Double, Vector Double)
layerNormBwd Matrix Double
x Vector Double
gamma Vector Double
beta Double
eps Matrix Double
dy
         in (Matrix Double
dx, (Vector Double
dgamma, Vector Double
dbeta))
    }

-- | Multi-head self-attention with a supplied attention mask.
--   Parameters are (Wq, Wk, Wv, Wo).  The mask is added to the raw QK^T
--   scores before softmax; use 'causalMaskF' for GPT and a zero matrix for
--   fully bidirectional BERT-style attention.
multiHeadAttentionP ::
  Int ->
  Int ->
  Double ->
  Matrix Double ->
  DiffP
    (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
    (Matrix Double)
    (Matrix Double)
multiHeadAttentionP :: Int
-> Int
-> Double
-> Matrix Double
-> DiffP
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
     (Matrix Double)
     (Matrix Double)
multiHeadAttentionP Int
nHead Int
seqLen Double
_eps Matrix Double
mask =
  TensorPrim
    { primForward :: (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
-> Matrix Double -> Matrix Double
primForward = \(Matrix Double
wq, Matrix Double
wk, Matrix Double
wv, Matrix Double
wo) Matrix Double
x ->
        let hDim :: Int
hDim = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
wq Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
nHead
            heads :: [Matrix Double]
heads =
              [ let wQh :: Matrix Double
wQh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wq Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wq) Int
hDim
                    wKh :: Matrix Double
wKh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wk Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wk) Int
hDim
                    wVh :: Matrix Double
wVh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wv Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wv) Int
hDim
                    q :: Matrix Double
q = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wQh
                    k :: Matrix Double
k = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wKh
                    v :: Matrix Double
v = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wVh
                    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 (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
hDim :: Double)) (Matrix Double
q Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
k)
                    probs :: Matrix Double
probs = Matrix Double -> Matrix Double
softmaxStableF (Matrix Double
scores Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
mask)
                 in Matrix Double
probs Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
v
              | Int
h <- [Int
0 .. Int
nHead Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
              ]
            ctx :: Matrix Double
ctx = [Matrix Double] -> Matrix Double
concatColsF [Matrix Double]
heads
         in Matrix Double
ctx Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wo,
      primBackward :: (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double,
    (Matrix Double, Matrix Double, Matrix Double, Matrix Double))
primBackward = \(Matrix Double
wq, Matrix Double
wk, Matrix Double
wv, Matrix Double
wo) Matrix Double
x Matrix Double
dOut ->
        let hDim :: Int
hDim = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
wq Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
nHead
            dk :: Double
dk = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
hDim :: Double
            -- Recompute forward intermediates
            an :: Matrix Double
an = Matrix Double
x
            headFwd :: Int
-> (Int, Matrix Double, Matrix Double, Matrix Double,
    Matrix Double, Matrix Double, Matrix Double, Matrix Double)
headFwd Int
h =
              let wQh :: Matrix Double
wQh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wq Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wq) Int
hDim
                  wKh :: Matrix Double
wKh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wk Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wk) Int
hDim
                  wVh :: Matrix Double
wVh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wv Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wv) Int
hDim
                  q :: Matrix Double
q = Matrix Double
an Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wQh
                  k :: Matrix Double
k = Matrix Double
an Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wKh
                  v :: Matrix Double
v = Matrix Double
an Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wVh
                  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
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
k)
                  probs :: Matrix Double
probs = Matrix Double -> Matrix Double
softmaxStableF (Matrix Double
scores Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
mask)
               in (Int
h, Matrix Double
wQh, Matrix Double
wKh, Matrix Double
wVh, Matrix Double
q, Matrix Double
k, Matrix Double
v, Matrix Double
probs)
            headStates :: [(Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
  Matrix Double, Matrix Double, Matrix Double)]
headStates = (Int
 -> (Int, Matrix Double, Matrix Double, Matrix Double,
     Matrix Double, Matrix Double, Matrix Double, Matrix Double))
-> [Int]
-> [(Int, Matrix Double, Matrix Double, Matrix Double,
     Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
forall a b. (a -> b) -> [a] -> [b]
map Int
-> (Int, Matrix Double, Matrix Double, Matrix Double,
    Matrix Double, Matrix Double, Matrix Double, Matrix Double)
headFwd [Int
0 .. Int
nHead Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ctx :: Matrix Double
ctx = [Matrix Double] -> Matrix Double
concatColsF [Matrix Double
probs Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
v | (Int
_, Matrix Double
_, Matrix Double
_, Matrix Double
_, Matrix Double
_, Matrix Double
_, Matrix Double
v, Matrix Double
probs) <- [(Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
  Matrix Double, Matrix Double, Matrix Double)]
headStates]
            -- Backward through output projection
            (Matrix Double
gradCtx, Matrix Double
gWo, Vector Double
_) = Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double, Vector Double)
forall {t}.
Numeric t =>
Matrix t -> Matrix t -> Matrix t -> (Matrix t, Matrix t, Vector t)
linearBwd Matrix Double
ctx Matrix Double
wo Matrix Double
dOut
            -- Per-head backward
            headBwd :: (Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
 Matrix Double, Matrix Double, Matrix Double)
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
headBwd (Int
h, Matrix Double
wQh, Matrix Double
wKh, Matrix Double
wVh, Matrix Double
q, Matrix Double
k, Matrix Double
v, Matrix Double
probs) =
              let goh :: Matrix Double
goh = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
gradCtx Int
0 (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim) Int
seqLen Int
hDim
                  gradProbs :: Matrix Double
gradProbs = Matrix Double
goh Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
v
                  gradV :: Matrix Double
gradV = Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
probs Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
goh
                  gradScores :: Matrix Double
gradScores = Matrix Double -> Matrix Double -> Matrix Double
forall {t}. (Element t, Num t) => Matrix t -> Matrix t -> Matrix t
softmaxBwd Matrix Double
probs Matrix Double
gradProbs
                  gradScoresS :: Matrix Double
gradScoresS = 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
gradScores
                  gradQ :: Matrix Double
gradQ = Matrix Double
gradScoresS Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
k
                  gradK :: Matrix Double
gradK = Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
gradScoresS Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
q
                  (Matrix Double
gradXq, Matrix Double
gWQ, Vector Double
_) = Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double, Vector Double)
forall {t}.
Numeric t =>
Matrix t -> Matrix t -> Matrix t -> (Matrix t, Matrix t, Vector t)
linearBwd Matrix Double
an Matrix Double
wQh Matrix Double
gradQ
                  (Matrix Double
gradXk, Matrix Double
gWK, Vector Double
_) = Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double, Vector Double)
forall {t}.
Numeric t =>
Matrix t -> Matrix t -> Matrix t -> (Matrix t, Matrix t, Vector t)
linearBwd Matrix Double
an Matrix Double
wKh Matrix Double
gradK
                  (Matrix Double
gradXv, Matrix Double
gWV, Vector Double
_) = Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double, Vector Double)
forall {t}.
Numeric t =>
Matrix t -> Matrix t -> Matrix t -> (Matrix t, Matrix t, Vector t)
linearBwd Matrix Double
an Matrix Double
wVh Matrix Double
gradV
               in (Matrix Double
gradXq Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
gradXk Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
gradXv, Matrix Double
gWQ, Matrix Double
gWK, Matrix Double
gWV)
            headGrads :: [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
headGrads = ((Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
  Matrix Double, Matrix Double, Matrix Double)
 -> (Matrix Double, Matrix Double, Matrix Double, Matrix Double))
-> [(Int, Matrix Double, Matrix Double, Matrix Double,
     Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
-> [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
 Matrix Double, Matrix Double, Matrix Double)
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
headBwd [(Int, Matrix Double, Matrix Double, Matrix Double, Matrix Double,
  Matrix Double, Matrix Double, Matrix Double)]
headStates
            gradX :: Matrix Double
gradX = (Matrix Double -> Matrix Double -> Matrix Double)
-> [Matrix Double] -> Matrix Double
forall a. HasCallStack => (a -> a -> a) -> [a] -> a
foldl1' Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
(+) [Matrix Double
gx | (Matrix Double
gx, Matrix Double
_, Matrix Double
_, Matrix Double
_) <- [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
headGrads]
            gWqFull :: Matrix Double
gWqFull = Matrix Double -> [(Int, Matrix Double)] -> Matrix Double
accumulateCols Matrix Double
wq [(Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim, Matrix Double
gwq) | (Int
h, (Matrix Double
_, Matrix Double
gwq, Matrix Double
_, Matrix Double
_)) <- [Int]
-> [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
-> [(Int,
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
headGrads]
            gWkFull :: Matrix Double
gWkFull = Matrix Double -> [(Int, Matrix Double)] -> Matrix Double
accumulateCols Matrix Double
wk [(Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim, Matrix Double
gwk) | (Int
h, (Matrix Double
_, Matrix Double
_, Matrix Double
gwk, Matrix Double
_)) <- [Int]
-> [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
-> [(Int,
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
headGrads]
            gWvFull :: Matrix Double
gWvFull = Matrix Double -> [(Int, Matrix Double)] -> Matrix Double
accumulateCols Matrix Double
wv [(Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDim, Matrix Double
gwv) | (Int
h, (Matrix Double
_, Matrix Double
_, Matrix Double
_, Matrix Double
gwv)) <- [Int]
-> [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
-> [(Int,
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Matrix Double, Matrix Double, Matrix Double, Matrix Double)]
headGrads]
         in (Matrix Double
gradX, (Matrix Double
gWqFull, Matrix Double
gWkFull, Matrix Double
gWvFull, Matrix Double
gWo))
    }
  where
    softmaxBwd :: Matrix t -> Matrix t -> Matrix t
softmaxBwd Matrix t
probs Matrix t
gradOut =
      [Vector t] -> Matrix t
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector t] -> Matrix t) -> [Vector t] -> Matrix t
forall a b. (a -> b) -> a -> b
$ (Vector t -> Vector t -> Vector t)
-> [Vector t] -> [Vector t] -> [Vector t]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Vector t -> Vector t -> Vector t
forall {a}. (Storable a, Num a) => Vector a -> Vector a -> Vector a
softmaxRowBwd (Matrix t -> [Vector t]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix t
probs) (Matrix t -> [Vector t]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix t
gradOut)
      where
        softmaxRowBwd :: Vector a -> Vector a -> Vector a
softmaxRowBwd Vector a
p Vector a
go =
          let pv :: [a]
pv = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
p; gov :: [a]
gov = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
go; dot :: a
dot = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
pv [a]
gov)
           in [a] -> Vector a
forall a. Storable a => [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
pi_ a
goi -> a
pi_ a -> a -> a
forall a. Num a => a -> a -> a
* (a
goi a -> a -> a
forall a. Num a => a -> a -> a
- a
dot)) [a]
pv [a]
gov
    linearBwd :: Matrix t -> Matrix t -> Matrix t -> (Matrix t, Matrix t, Vector t)
linearBwd Matrix t
x Matrix t
w Matrix t
gradY =
      let gradX :: Matrix t
gradX = Matrix t
gradY Matrix t -> Matrix t -> Matrix t
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix t -> Matrix t
forall m mt. Transposable m mt => m -> mt
tr Matrix t
w
          gradW :: Matrix t
gradW = Matrix t -> Matrix t
forall m mt. Transposable m mt => m -> mt
tr Matrix t
x Matrix t -> Matrix t -> Matrix t
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix t
gradY
          gradB :: Vector t
gradB = [t] -> Vector t
forall a. Storable a => [a] -> Vector a
fromList [Vector t -> t
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector t
col | Vector t
col <- Matrix t -> [Vector t]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix t
gradY]
       in (Matrix t
gradX, Matrix t
gradW, Vector t
gradB)

----------------------------------------------------------------------
-- Transformer block
----------------------------------------------------------------------

-- | Flat parameter bundle for a transformer block.
data BlockParams = BlockParams
  { BlockParams -> Matrix Double
bpAttnWq, BlockParams -> Matrix Double
bpAttnWk, BlockParams -> Matrix Double
bpAttnWv, BlockParams -> Matrix Double
bpAttnWo :: Matrix Double,
    BlockParams -> Vector Double
bpAttnLnGamma, BlockParams -> Vector Double
bpAttnLnBeta :: Vector Double,
    BlockParams -> Matrix Double
bpFfnW1, BlockParams -> Matrix Double
bpFfnW2 :: Matrix Double,
    BlockParams -> Vector Double
bpFfnB1, BlockParams -> Vector Double
bpFfnB2 :: Vector Double,
    BlockParams -> Vector Double
bpFfnLnGamma, BlockParams -> Vector Double
bpFfnLnBeta :: Vector Double
  }

blockParamsFromBlock :: TransformerBlock -> BlockParams
blockParamsFromBlock :: TransformerBlock -> BlockParams
blockParamsFromBlock TransformerBlock
tb =
  let ff :: FeedForward
ff = TransformerBlock -> FeedForward
tbFfn TransformerBlock
tb
   in BlockParams
        { bpAttnWq :: Matrix Double
bpAttnWq = TransformerBlock -> Matrix Double
tbAttnWq TransformerBlock
tb,
          bpAttnWk :: Matrix Double
bpAttnWk = TransformerBlock -> Matrix Double
tbAttnWk TransformerBlock
tb,
          bpAttnWv :: Matrix Double
bpAttnWv = TransformerBlock -> Matrix Double
tbAttnWv TransformerBlock
tb,
          bpAttnWo :: Matrix Double
bpAttnWo = TransformerBlock -> Matrix Double
tbAttnWo TransformerBlock
tb,
          bpAttnLnGamma :: Vector Double
bpAttnLnGamma = TransformerBlock -> Vector Double
tbAttnLnGamma TransformerBlock
tb,
          bpAttnLnBeta :: Vector Double
bpAttnLnBeta = TransformerBlock -> Vector Double
tbAttnLnBeta TransformerBlock
tb,
          bpFfnW1 :: Matrix Double
bpFfnW1 = FeedForward -> Matrix Double
ffW1 FeedForward
ff,
          bpFfnB1 :: Vector Double
bpFfnB1 = FeedForward -> Vector Double
ffB1 FeedForward
ff,
          bpFfnW2 :: Matrix Double
bpFfnW2 = FeedForward -> Matrix Double
ffW2 FeedForward
ff,
          bpFfnB2 :: Vector Double
bpFfnB2 = FeedForward -> Vector Double
ffB2 FeedForward
ff,
          bpFfnLnGamma :: Vector Double
bpFfnLnGamma = TransformerBlock -> Vector Double
tbFfnLnGamma TransformerBlock
tb,
          bpFfnLnBeta :: Vector Double
bpFfnLnBeta = TransformerBlock -> Vector Double
tbFfnLnBeta TransformerBlock
tb
        }

-- | Post-LN transformer block as a differentiable operation.
--
--   x -> Attention -> + -> LN -> FFN -> + -> LN -> y
--
-- The implementation is written directly against 'BlockParams' rather than
-- composed with '@.' so that the parameter type stays flat and readable.
-- The attention mask is supplied explicitly, so the same block can be used
-- for causal (GPT) or bidirectional (BERT) attention.
blockDiffP ::
  Int ->
  Int ->
  Double ->
  Matrix Double ->
  DiffP BlockParams (Matrix Double) (Matrix Double)
blockDiffP :: Int
-> Int
-> Double
-> Matrix Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
blockDiffP Int
nHead Int
seqLen Double
eps Matrix Double
mask =
  TensorPrim
    { primForward :: BlockParams -> Matrix Double -> Matrix Double
primForward = \BlockParams
p Matrix Double
x ->
        let attnOut :: Matrix Double
attnOut =
              DiffP
  (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
  (Matrix Double)
  (Matrix Double)
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
-> Matrix Double
-> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward
                (Int
-> Int
-> Double
-> Matrix Double
-> DiffP
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
     (Matrix Double)
     (Matrix Double)
multiHeadAttentionP Int
nHead Int
seqLen Double
eps Matrix Double
mask)
                (BlockParams -> Matrix Double
bpAttnWq BlockParams
p, BlockParams -> Matrix Double
bpAttnWk BlockParams
p, BlockParams -> Matrix Double
bpAttnWv BlockParams
p, BlockParams -> Matrix Double
bpAttnWo BlockParams
p)
                Matrix Double
x
            postAttn :: Matrix Double
postAttn = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
attnOut
            attnNorm :: Matrix Double
attnNorm = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
postAttn (BlockParams -> Vector Double
bpAttnLnGamma BlockParams
p) (BlockParams -> Vector Double
bpAttnLnBeta BlockParams
p) Double
eps
            ffnHidden :: Matrix Double
ffnHidden = DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double) -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW1 BlockParams
p, BlockParams -> Vector Double
bpFfnB1 BlockParams
p) Matrix Double
attnNorm
            ffnActivated :: Matrix Double
ffnActivated = Matrix Double -> Matrix Double
geluF Matrix Double
ffnHidden
            ffnOut :: Matrix Double
ffnOut = DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double) -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW2 BlockParams
p, BlockParams -> Vector Double
bpFfnB2 BlockParams
p) Matrix Double
ffnActivated
            ffnRes :: Matrix Double
ffnRes = Matrix Double
attnNorm Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
ffnOut
         in Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
ffnRes (BlockParams -> Vector Double
bpFfnLnGamma BlockParams
p) (BlockParams -> Vector Double
bpFfnLnBeta BlockParams
p) Double
eps,
      primBackward :: BlockParams
-> Matrix Double -> Matrix Double -> (Matrix Double, BlockParams)
primBackward = \BlockParams
p Matrix Double
x Matrix Double
dy ->
        -- Forward recompute
        let attnOut :: Matrix Double
attnOut =
              DiffP
  (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
  (Matrix Double)
  (Matrix Double)
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
-> Matrix Double
-> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward
                (Int
-> Int
-> Double
-> Matrix Double
-> DiffP
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
     (Matrix Double)
     (Matrix Double)
multiHeadAttentionP Int
nHead Int
seqLen Double
eps Matrix Double
mask)
                (BlockParams -> Matrix Double
bpAttnWq BlockParams
p, BlockParams -> Matrix Double
bpAttnWk BlockParams
p, BlockParams -> Matrix Double
bpAttnWv BlockParams
p, BlockParams -> Matrix Double
bpAttnWo BlockParams
p)
                Matrix Double
x
            postAttn :: Matrix Double
postAttn = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
attnOut
            attnNorm :: Matrix Double
attnNorm = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
postAttn (BlockParams -> Vector Double
bpAttnLnGamma BlockParams
p) (BlockParams -> Vector Double
bpAttnLnBeta BlockParams
p) Double
eps
            ffnHidden :: Matrix Double
ffnHidden = DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double) -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW1 BlockParams
p, BlockParams -> Vector Double
bpFfnB1 BlockParams
p) Matrix Double
attnNorm
            ffnActivated :: Matrix Double
ffnActivated = Matrix Double -> Matrix Double
geluF Matrix Double
ffnHidden
            ffnOut :: Matrix Double
ffnOut = DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double) -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW2 BlockParams
p, BlockParams -> Vector Double
bpFfnB2 BlockParams
p) Matrix Double
ffnActivated
            ffnRes :: Matrix Double
ffnRes = Matrix Double
attnNorm Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
ffnOut
            -- Backward through final LN
            (Matrix Double
dFfnRes, (Vector Double
dgFfnGamma, Vector Double
dgFfnBeta)) =
              DiffP
  (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Vector Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Vector Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward (Double
-> DiffP
     (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
layerNormP Double
eps) (BlockParams -> Vector Double
bpFfnLnGamma BlockParams
p, BlockParams -> Vector Double
bpFfnLnBeta BlockParams
p) Matrix Double
ffnRes Matrix Double
dy
            -- Backward through FFN residual: dFfnRes splits to attnNorm and FFN
            (Matrix Double
dAttnNorm2, (Matrix Double
dgFfnW2, Vector Double
dgFfnB2)) =
              DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Matrix Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW2 BlockParams
p, BlockParams -> Vector Double
bpFfnB2 BlockParams
p) Matrix Double
ffnActivated Matrix Double
dFfnRes
            (Matrix Double
dFfnHidden, ()) =
              (Matrix Double, Matrix Double -> (Matrix Double, ()))
-> Matrix Double -> (Matrix Double, ())
forall a b. (a, b) -> b
snd (DiffP () (Matrix Double) (Matrix Double)
-> ()
-> Matrix Double
-> (Matrix Double, Matrix Double -> (Matrix Double, ()))
forall p a b. DiffP p a b -> p -> a -> (b, b -> (a, p))
ADP.runDiffP DiffP () (Matrix Double) (Matrix Double)
geluP () Matrix Double
ffnHidden) Matrix Double
dAttnNorm2
            (Matrix Double
dAttnNorm1, (Matrix Double
dgFfnW1, Vector Double
dgFfnB1)) =
              DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Matrix Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (BlockParams -> Matrix Double
bpFfnW1 BlockParams
p, BlockParams -> Vector Double
bpFfnB1 BlockParams
p) Matrix Double
attnNorm Matrix Double
dFfnHidden
            dAttnNorm :: Matrix Double
dAttnNorm = Matrix Double
dFfnRes Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
dAttnNorm1
            -- Backward through attn LN
            (Matrix Double
dPostAttn, (Vector Double
dgAttnGamma, Vector Double
dgAttnBeta)) =
              DiffP
  (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Vector Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Vector Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward (Double
-> DiffP
     (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
layerNormP Double
eps) (BlockParams -> Vector Double
bpAttnLnGamma BlockParams
p, BlockParams -> Vector Double
bpAttnLnBeta BlockParams
p) Matrix Double
postAttn Matrix Double
dAttnNorm
            -- Backward through attn residual: dPostAttn splits to x and attnOut
            (Matrix Double
dXfromAttn, (Matrix Double
dgWq, Matrix Double
dgWk, Matrix Double
dgWv, Matrix Double
dgWo)) =
              DiffP
  (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
  (Matrix Double)
  (Matrix Double)
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double,
    (Matrix Double, Matrix Double, Matrix Double, Matrix Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward
                (Int
-> Int
-> Double
-> Matrix Double
-> DiffP
     (Matrix Double, Matrix Double, Matrix Double, Matrix Double)
     (Matrix Double)
     (Matrix Double)
multiHeadAttentionP Int
nHead Int
seqLen Double
eps Matrix Double
mask)
                (BlockParams -> Matrix Double
bpAttnWq BlockParams
p, BlockParams -> Matrix Double
bpAttnWk BlockParams
p, BlockParams -> Matrix Double
bpAttnWv BlockParams
p, BlockParams -> Matrix Double
bpAttnWo BlockParams
p)
                Matrix Double
x
                Matrix Double
dPostAttn
            dX :: Matrix Double
dX = Matrix Double
dPostAttn Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
dXfromAttn
         in ( Matrix Double
dX,
              BlockParams
                { bpAttnWq :: Matrix Double
bpAttnWq = Matrix Double
dgWq,
                  bpAttnWk :: Matrix Double
bpAttnWk = Matrix Double
dgWk,
                  bpAttnWv :: Matrix Double
bpAttnWv = Matrix Double
dgWv,
                  bpAttnWo :: Matrix Double
bpAttnWo = Matrix Double
dgWo,
                  bpAttnLnGamma :: Vector Double
bpAttnLnGamma = Vector Double
dgAttnGamma,
                  bpAttnLnBeta :: Vector Double
bpAttnLnBeta = Vector Double
dgAttnBeta,
                  bpFfnW1 :: Matrix Double
bpFfnW1 = Matrix Double
dgFfnW1,
                  bpFfnB1 :: Vector Double
bpFfnB1 = Vector Double
dgFfnB1,
                  bpFfnW2 :: Matrix Double
bpFfnW2 = Matrix Double
dgFfnW2,
                  bpFfnB2 :: Vector Double
bpFfnB2 = Vector Double
dgFfnB2,
                  bpFfnLnGamma :: Vector Double
bpFfnLnGamma = Vector Double
dgFfnGamma,
                  bpFfnLnBeta :: Vector Double
bpFfnLnBeta = Vector Double
dgFfnBeta
                }
            )
    }

-- | Causal transformer block (GPT-style).
gptBlockDiffP :: Int -> Int -> Double -> DiffP BlockParams (Matrix Double) (Matrix Double)
gptBlockDiffP :: Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
gptBlockDiffP Int
nHead Int
seqLen Double
eps = Int
-> Int
-> Double
-> Matrix Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
blockDiffP Int
nHead Int
seqLen Double
eps (Int -> Matrix Double
causalMaskF Int
seqLen)

-- | Bidirectional transformer block (BERT-style).  Attention can attend to
-- all positions in the sequence.
bertBlockDiffP :: Int -> Int -> Double -> DiffP BlockParams (Matrix Double) (Matrix Double)
bertBlockDiffP :: Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
bertBlockDiffP Int
nHead Int
seqLen Double
eps = Int
-> Int
-> Double
-> Matrix Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
blockDiffP Int
nHead Int
seqLen Double
eps (Int -> Int -> Matrix Double
zeroMatrix Int
seqLen Int
seqLen)

----------------------------------------------------------------------
-- Full model bodies
----------------------------------------------------------------------

-- | Flat parameter bundle for the whole transformer body (excluding
--   embeddings, which are handled separately).
data GptParams = GptParams
  { GptParams -> [BlockParams]
gpBlocks :: [BlockParams],
    GptParams -> Vector Double
gpLnGamma, GptParams -> Vector Double
gpLnBeta :: Vector Double,
    GptParams -> Matrix Double
gpHead :: Matrix Double,
    GptParams -> Vector Double
gpHeadB :: Vector Double
  }

gptParamsFromModel :: Gpt -> GptParams
gptParamsFromModel :: Gpt -> GptParams
gptParamsFromModel Gpt
m =
  GptParams
    { gpBlocks :: [BlockParams]
gpBlocks = (TransformerBlock -> BlockParams)
-> [TransformerBlock] -> [BlockParams]
forall a b. (a -> b) -> [a] -> [b]
map TransformerBlock -> BlockParams
blockParamsFromBlock (Gpt -> [TransformerBlock]
gptBlocks Gpt
m),
      gpLnGamma :: Vector Double
gpLnGamma = Gpt -> Vector Double
gptLnGamma Gpt
m,
      gpLnBeta :: Vector Double
gpLnBeta = Gpt -> Vector Double
gptLnBeta Gpt
m,
      gpHead :: Matrix Double
gpHead = Gpt -> Matrix Double
gptHead Gpt
m,
      gpHeadB :: Vector Double
gpHeadB = Gpt -> Vector Double
gptHeadB Gpt
m
    }

-- | Generic transformer body.  The supplied block constructor chooses the
--   attention mask (causal for GPT, bidirectional for BERT).
bodyDiffP ::
  GptConfig ->
  Int ->
  Double ->
  (Int -> Int -> Double -> DiffP BlockParams (Matrix Double) (Matrix Double)) ->
  DiffP GptParams (Matrix Double) (Matrix Double)
bodyDiffP :: GptConfig
-> Int
-> Double
-> (Int
    -> Int
    -> Double
    -> DiffP BlockParams (Matrix Double) (Matrix Double))
-> DiffP GptParams (Matrix Double) (Matrix Double)
bodyDiffP GptConfig
cfg Int
seqLen Double
eps Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
blockCtor =
  TensorPrim
    { primForward :: GptParams -> Matrix Double -> Matrix Double
primForward = \GptParams
p Matrix Double
x ->
        let xBlocks :: Matrix Double
xBlocks = (Matrix Double -> BlockParams -> Matrix Double)
-> Matrix Double -> [BlockParams] -> 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' Matrix Double -> BlockParams -> Matrix Double
applyBlock Matrix Double
x (GptParams -> [BlockParams]
gpBlocks GptParams
p)
            xFinalNorm :: Matrix Double
xFinalNorm = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
xBlocks (GptParams -> Vector Double
gpLnGamma GptParams
p) (GptParams -> Vector Double
gpLnBeta GptParams
p) Double
eps
         in Matrix Double
xFinalNorm Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> GptParams -> Matrix Double
gpHead GptParams
p Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Vector Double -> Int -> Matrix Double
broadcastBias (GptParams -> Vector Double
gpHeadB GptParams
p) Int
seqLen,
      primBackward :: GptParams
-> Matrix Double -> Matrix Double -> (Matrix Double, GptParams)
primBackward = \GptParams
p Matrix Double
x Matrix Double
dy ->
        let fwdStack :: [Matrix Double]
fwdStack = ([Matrix Double] -> BlockParams -> [Matrix Double])
-> [Matrix Double] -> [BlockParams] -> [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' (\[Matrix Double]
acc BlockParams
bp -> [Matrix Double]
acc [Matrix Double] -> [Matrix Double] -> [Matrix Double]
forall a. [a] -> [a] -> [a]
++ [Matrix Double -> BlockParams -> Matrix Double
applyBlock ([Matrix Double] -> Matrix Double
forall a. HasCallStack => [a] -> a
last [Matrix Double]
acc) BlockParams
bp]) [Matrix Double
x] (GptParams -> [BlockParams]
gpBlocks GptParams
p)
            xBlocks :: Matrix Double
xBlocks = [Matrix Double] -> Matrix Double
forall a. HasCallStack => [a] -> a
last [Matrix Double]
fwdStack
            xFinalNorm :: Matrix Double
xFinalNorm = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
xBlocks (GptParams -> Vector Double
gpLnGamma GptParams
p) (GptParams -> Vector Double
gpLnBeta GptParams
p) Double
eps
            -- Backward through output projection
            (Matrix Double
gradXFinalNorm, (Matrix Double
gHead, Vector Double
gHeadB)) =
              DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Matrix Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Matrix Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP
  (Matrix Double, Vector Double) (Matrix Double) (Matrix Double)
linearP (GptParams -> Matrix Double
gpHead GptParams
p, GptParams -> Vector Double
gpHeadB GptParams
p) Matrix Double
xFinalNorm Matrix Double
dy
            -- Backward through final layer norm
            (Matrix Double
gradXBlocks, (Vector Double
gLnGamma, Vector Double
gLnBeta)) =
              DiffP
  (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
-> (Vector Double, Vector Double)
-> Matrix Double
-> Matrix Double
-> (Matrix Double, (Vector Double, Vector Double))
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward (Double
-> DiffP
     (Vector Double, Vector Double) (Matrix Double) (Matrix Double)
layerNormP Double
eps) (GptParams -> Vector Double
gpLnGamma GptParams
p, GptParams -> Vector Double
gpLnBeta GptParams
p) Matrix Double
xBlocks Matrix Double
gradXFinalNorm
            -- Backward through blocks
            xPrevs :: [Matrix Double]
xPrevs = case [Matrix Double] -> [Matrix Double]
forall a. [a] -> [a]
reverse [Matrix Double]
fwdStack of (Matrix Double
_ : [Matrix Double]
xs) -> [Matrix Double]
xs; [] -> []
            (Matrix Double
gradX0, [BlockParams]
blockGradsRev) =
              Matrix Double
-> [BlockParams]
-> [Matrix Double]
-> (Matrix Double, [BlockParams])
goBwd Matrix Double
gradXBlocks ([BlockParams] -> [BlockParams]
forall a. [a] -> [a]
reverse (GptParams -> [BlockParams]
gpBlocks GptParams
p)) [Matrix Double]
xPrevs
         in ( Matrix Double
gradX0,
              GptParams
                { gpBlocks :: [BlockParams]
gpBlocks = [BlockParams] -> [BlockParams]
forall a. [a] -> [a]
reverse [BlockParams]
blockGradsRev,
                  gpLnGamma :: Vector Double
gpLnGamma = Vector Double
gLnGamma,
                  gpLnBeta :: Vector Double
gpLnBeta = Vector Double
gLnBeta,
                  gpHead :: Matrix Double
gpHead = Matrix Double
gHead,
                  gpHeadB :: Vector Double
gpHeadB = Vector Double
gHeadB
                }
            )
    }
  where
    nHead :: Int
nHead = GptConfig -> Int
gptNHead GptConfig
cfg
    blockP :: DiffP BlockParams (Matrix Double) (Matrix Double)
blockP = Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
blockCtor Int
nHead Int
seqLen Double
eps
    applyBlock :: Matrix Double -> BlockParams -> Matrix Double
applyBlock Matrix Double
xIn BlockParams
bp = DiffP BlockParams (Matrix Double) (Matrix Double)
-> BlockParams -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP BlockParams (Matrix Double) (Matrix Double)
blockP BlockParams
bp Matrix Double
xIn
    goBwd :: Matrix Double
-> [BlockParams]
-> [Matrix Double]
-> (Matrix Double, [BlockParams])
goBwd Matrix Double
gradX [] [Matrix Double]
_ = (Matrix Double
gradX, [])
    goBwd Matrix Double
gradX (BlockParams
bp : [BlockParams]
bps') (Matrix Double
xPrev : [Matrix Double]
xs) =
      let (Matrix Double
dx, BlockParams
dbp) = DiffP BlockParams (Matrix Double) (Matrix Double)
-> BlockParams
-> Matrix Double
-> Matrix Double
-> (Matrix Double, BlockParams)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP BlockParams (Matrix Double) (Matrix Double)
blockP BlockParams
bp Matrix Double
xPrev Matrix Double
gradX
          (Matrix Double
dxFinal, [BlockParams]
dbps) = Matrix Double
-> [BlockParams]
-> [Matrix Double]
-> (Matrix Double, [BlockParams])
goBwd Matrix Double
dx [BlockParams]
bps' [Matrix Double]
xs
       in (Matrix Double
dxFinal, BlockParams
dbp BlockParams -> [BlockParams] -> [BlockParams]
forall a. a -> [a] -> [a]
: [BlockParams]
dbps)
    goBwd Matrix Double
_ [BlockParams]
_ [Matrix Double]
_ = [Char] -> (Matrix Double, [BlockParams])
forall a. HasCallStack => [Char] -> a
error [Char]
"bodyDiffP: mismatched block stack"

-- | Full GPT model as a differentiable operation.
--
-- Input is the already-embedded token matrix @x0@; output is logits.
-- Embeddings are kept outside because the lookup is not differentiable wrt
-- token IDs, only wrt the embedding matrices.
gptDiffP ::
  GptConfig ->
  Int ->
  Double ->
  DiffP GptParams (Matrix Double) (Matrix Double)
gptDiffP :: GptConfig
-> Int -> Double -> DiffP GptParams (Matrix Double) (Matrix Double)
gptDiffP GptConfig
cfg Int
seqLen Double
eps = GptConfig
-> Int
-> Double
-> (Int
    -> Int
    -> Double
    -> DiffP BlockParams (Matrix Double) (Matrix Double))
-> DiffP GptParams (Matrix Double) (Matrix Double)
bodyDiffP GptConfig
cfg Int
seqLen Double
eps Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
gptBlockDiffP

-- | Full BERT-style bidirectional model as a differentiable operation.
--   The architecture is identical to GPT except that attention is not
--   causally masked, so every position can attend to every other position.
bertDiffP ::
  GptConfig ->
  Int ->
  Double ->
  DiffP GptParams (Matrix Double) (Matrix Double)
bertDiffP :: GptConfig
-> Int -> Double -> DiffP GptParams (Matrix Double) (Matrix Double)
bertDiffP GptConfig
cfg Int
seqLen Double
eps = GptConfig
-> Int
-> Double
-> (Int
    -> Int
    -> Double
    -> DiffP BlockParams (Matrix Double) (Matrix Double))
-> DiffP GptParams (Matrix Double) (Matrix Double)
bodyDiffP GptConfig
cfg Int
seqLen Double
eps Int
-> Int
-> Double
-> DiffP BlockParams (Matrix Double) (Matrix Double)
bertBlockDiffP

----------------------------------------------------------------------
-- Forward helpers (shared with Backprop)
----------------------------------------------------------------------

geluF :: Matrix Double -> Matrix Double
geluF :: Matrix Double -> Matrix Double
geluF = (Double -> Double) -> Matrix Double -> Matrix Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap Double -> Double
forall a. Floating a => a -> a
f
  where
    f :: a -> a
f a
v = let z :: a
z = a
1.59577 a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
0.044715 a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v) in a
v a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Floating a => a -> a
exp (-a
z))

-- | GELU backward helper (used by the migrated 'geluP').
geluBwd :: Matrix Double -> Matrix Double -> Matrix Double
geluBwd :: Matrix Double -> Matrix Double -> Matrix Double
geluBwd Matrix Double
x Matrix Double
gradY =
  let xs :: [Double]
xs = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
x)
      gs :: [Double]
gs = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
gradY)
      deriv :: a -> a
deriv a
v =
        let a :: a
a = a
1.59577
            b :: a
b = a
0.044715
            z :: a
z = a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v)
            phi :: a
phi = a
1 a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Floating a => a -> a
exp (-a
z))
            phi' :: a
phi' = a
phi a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
phi) a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v)
         in a
phi a -> a -> a
forall a. Num a => a -> a -> a
+ a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
phi'
   in Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x) ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
gs ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Floating a => a -> a
deriv [Double]
xs)))

softmaxStableF :: Matrix Double -> Matrix Double
softmaxStableF :: Matrix Double -> Matrix Double
softmaxStableF Matrix Double
m =
  [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows
    [ let mx :: Double
mx = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
maxElement Vector Double
row; s :: Vector Double
s = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
x -> Double -> Double
forall a. Floating a => a -> a
exp (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx)) Vector Double
row
       in Double -> Vector Double -> Vector 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
/ Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
s) Vector Double
s
    | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
m
    ]

layerNormF :: Matrix Double -> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF :: Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNormF Matrix Double
x Vector Double
gamma Vector Double
beta Double
eps =
  let d :: Double
d = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x)
      mu :: Vector Double
mu = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
row Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x]
      xm :: Matrix Double
xm = Matrix Double
x 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.asColumn Vector Double
mu
      var :: Vector Double
var = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
row Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
row) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm]
      invStd :: Vector Double
invStd = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
v -> Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)) Vector Double
var
      xHat :: Matrix Double
xHat = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector Double] -> Matrix Double)
-> [Vector Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Double -> Vector Double)
-> [Vector Double] -> [Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Double -> Vector Double -> Vector Double)
-> Vector Double -> Double -> Vector Double
forall a b c. (a -> b -> c) -> b -> a -> c
flip Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
invStd)
   in [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows
        [ [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
v Double
g Double
b -> Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
g Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
row) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
gamma) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
beta)
        | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xHat
        ]

layerNormBwd ::
  Matrix Double ->
  Vector Double ->
  Vector Double ->
  Double ->
  Matrix Double ->
  (Matrix Double, Vector Double, Vector Double)
layerNormBwd :: Matrix Double
-> Vector Double
-> Vector Double
-> Double
-> Matrix Double
-> (Matrix Double, Vector Double, Vector Double)
layerNormBwd Matrix Double
x Vector Double
gamma Vector Double
_beta Double
eps Matrix Double
gradY =
  let mu :: Vector Double
mu = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
row Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x]
      xm :: Matrix Double
xm = Matrix Double
x 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.asColumn Vector Double
mu
      var :: Vector Double
var = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
row Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
row) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm]
      invStd :: Vector Double
invStd = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
v -> Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)) Vector Double
var
      xHat :: Matrix Double
xHat = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector Double] -> Matrix Double)
-> [Vector Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Double -> Vector Double)
-> [Vector Double] -> [Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Double -> Vector Double -> Vector Double)
-> Vector Double -> Double -> Vector Double
forall a b c. (a -> b -> c) -> b -> a -> c
flip Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
invStd)
      gammaList :: [Double]
gammaList = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
gamma
      gradXHat :: Matrix Double
gradXHat = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [[Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
go) [Double]
gammaList | Vector Double
go <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradY]
      gGamma :: Vector Double
gGamma = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
go Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
xh) | (Vector Double
go, Vector Double
xh) <- [Vector Double]
-> [Vector Double] -> [(Vector Double, Vector Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
gradY) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
xHat)]
      gBeta :: Vector Double
gBeta = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
go | Vector Double
go <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
gradY]
      gradX :: Matrix Double
gradX =
        [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows
          [ Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRow
              (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradXHat [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
invStd [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
var [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
          | Int
i <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
          ]
   in (Matrix Double
gradX, Vector Double
gGamma, Vector Double
gBeta)
  where
    d :: Double
d = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x)
    features :: Int
features = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x
    gradRow :: Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRow Vector Double
gxh Vector Double
xi_minus_mu Double
is Double
v =
      let gInvStd :: Double
gInvStd = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
gxh Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
xi_minus_mu)
          gVar :: Double
gVar = Double
gInvStd Double -> Double -> Double
forall a. Num a => a -> a -> a
* (-Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
1.5)
          gXm :: Vector Double
gXm = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
is Vector Double
gxh Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
gVar Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d) Vector Double
xi_minus_mu
          gMu :: Double
gMu = -Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
gXm
       in Vector Double
gXm Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
features (Double
gMu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d))

concatColsF :: [Matrix Double] -> Matrix Double
concatColsF :: [Matrix Double] -> Matrix Double
concatColsF [Matrix Double
a] = Matrix Double
a
concatColsF (Matrix Double
a : [Matrix Double]
as) = Matrix Double
a Matrix Double -> Matrix Double -> Matrix Double
forall t. Element t => Matrix t -> Matrix t -> Matrix t
||| [Matrix Double] -> Matrix Double
concatColsF [Matrix Double]
as
concatColsF [] = [Char] -> Matrix Double
forall a. HasCallStack => [Char] -> a
error [Char]
"no heads"

broadcastBias :: Vector Double -> Int -> Matrix Double
broadcastBias :: Vector Double -> Int -> Matrix Double
broadcastBias Vector Double
b Int
n = Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape (Vector Double -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
LA.size Vector Double
b) ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate Int
n (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
b))))

subMatrixW :: Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW :: Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
m Int
r Int
c Int
rows' Int
cols' = (Int, Int) -> (Int, Int) -> Matrix Double -> Matrix Double
forall a.
Element a =>
(Int, Int) -> (Int, Int) -> Matrix a -> Matrix a
LA.subMatrix (Int
r, Int
c) (Int
rows', Int
cols') Matrix Double
m

causalMaskF :: Int -> Matrix Double
causalMaskF :: Int -> Matrix Double
causalMaskF 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]]

accumulateCols :: Matrix Double -> [(Int, Matrix Double)] -> Matrix Double
accumulateCols :: Matrix Double -> [(Int, Matrix Double)] -> Matrix Double
accumulateCols Matrix Double
template [(Int, Matrix Double)]
patches =
  let full :: Matrix Double
full = Int -> Int -> Matrix Double
zeroMatrix (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
template) (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
template)
   in (Matrix Double -> (Int, Matrix Double) -> Matrix Double)
-> Matrix Double -> [(Int, Matrix 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'
        ( \Matrix Double
acc (Int
offset, Matrix Double
patch) ->
            Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape
              (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc)
              ( [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList
                  [ if Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
offset Bool -> Bool -> Bool
&& Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
patch
                      then
                        (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
acc) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c))
                          Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
patch) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
patch Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offset)))
                      else Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
acc) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c)
                  | Int
r <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1],
                    Int
c <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                  ]
              )
        )
        Matrix Double
full
        [(Int, Matrix Double)]
patches

zeroMatrix :: Int -> Int -> Matrix Double
zeroMatrix :: Int -> Int -> Matrix Double
zeroMatrix Int
r Int
c = Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape Int
c ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c) Double
0))