-- SPDX-License-Identifier: MPL-2.0

{- | Row-major CPU layers. Convolution/pooling layout is N,C,spatial (1–3
spatial axes), weights are Cout,Cin/groups,kernel, or Cin,Cout/groups,kernel
for transpose convolution. Padding is symmetric and zero-filled. Attention
layout is sequence,feature; True in a mask means allowed. Fully masked rows
produce exact zero outputs and weights. No ambient training state exists.
-}
module HMath.DeepLearning.Layers (
    Training (..),
    dropout,
    spatialDropout,
    BatchNormState (..),
    batchNorm,
    layerNorm,
    instanceNorm,
    groupNorm,
    rmsNorm,
    weightNorm,
    Initializer (..),
    initializeWeights,
    embedding,
    dense,
    denseAD,
    ConvOptions (..),
    convolution,
    transposedConvolution,
    depthwiseConvolution,
    convolutionAD,
    PoolMode (..),
    pooling,
    adaptivePooling,
    globalPooling,
    scaledDotProductAttention,
    multiheadAttention,
    causalMask,
    rotaryEmbedding,
    sinusoidalEncoding,
    alibiBias,
    KVCache,
    emptyCache,
    appendCache,
    cacheValues,
    transformerFeedForward,
    layerNormAD,
    rmsNormAD,
    scaledDotProductAttentionAD,
) where

import Control.Monad (foldM)
import qualified Data.Foldable as Foldable (foldl')
import Data.List (nub)
import qualified HMath.Autodiff as AD
import HMath.DeepLearning.Core (Activation, activation, softmax, softmaxExpr)
import HMath.Error
import HMath.Random (RNG, normal, sampleN, uniform)
import qualified HMath.Tensor as T

{- | Layer mode: Training updates/uses batch statistics; Inference uses retained running
statistics.
-}
data Training = Training | Inference deriving (Training -> Training -> Bool
(Training -> Training -> Bool)
-> (Training -> Training -> Bool) -> Eq Training
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Training -> Training -> Bool
== :: Training -> Training -> Bool
$c/= :: Training -> Training -> Bool
/= :: Training -> Training -> Bool
Eq, Int -> Training -> ShowS
[Training] -> ShowS
Training -> String
(Int -> Training -> ShowS)
-> (Training -> String) -> ([Training] -> ShowS) -> Show Training
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Training -> ShowS
showsPrec :: Int -> Training -> ShowS
$cshow :: Training -> String
show :: Training -> String
$cshowList :: [Training] -> ShowS
showList :: [Training] -> ShowS
Show)

finite :: [Double] -> Either MathError ()
finite :: [Double] -> Either MathError ()
finite [Double]
xs = if (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)) [Double]
xs then () -> Either MathError ()
forall a b. b -> Either a b
Right () else MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
NonFiniteResult
validTensor :: T.Tensor Double -> Either MathError ()
validTensor :: Tensor Double -> Either MathError ()
validTensor = [Double] -> Either MathError ()
finite ([Double] -> Either MathError ())
-> (Tensor Double -> [Double])
-> Tensor Double
-> Either MathError ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues
indices :: [Int] -> [[Int]]
indices :: [Int] -> [[Int]]
indices [] = [[]]
indices (Int
d : [Int]
ds) = [Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
is | Int
i <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
is <- [Int] -> [[Int]]
indices [Int]
ds]
flat :: [Int] -> [Int] -> Int
flat :: [Int] -> [Int] -> Int
flat [Int]
shape [Int]
ix = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
ix (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
shape)))
get :: T.Tensor a -> [Int] -> a
get :: forall a. Tensor a -> [Int] -> a
get Tensor a
t [Int]
ix = Tensor a -> [a]
forall a. Tensor a -> [a]
T.tensorValues Tensor a
t [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! [Int] -> [Int] -> Int
flat (Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
t) [Int]
ix
checkedMake :: [Int] -> [Double] -> Either MathError (T.Tensor Double)
checkedMake :: [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake [Int]
shape [Double]
values = do [Double] -> Either MathError ()
finite [Double]
values; [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int]
shape [Double]
values

{- | Inverted dropout. Inference consumes no random numbers. Return a separate
multiplicative mask so its exact graph can be replayed by backpropagation.
-}
dropout ::
    Training ->
    Double ->
    RNG ->
    T.Tensor Double ->
    Either MathError (T.Tensor Double, T.Tensor Double, RNG)
dropout :: Training
-> Double
-> RNG
-> Tensor Double
-> Either MathError (Tensor Double, Tensor Double, RNG)
dropout Training
mode Double
probability RNG
rng Tensor Double
input = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    [Double] -> Either MathError ()
finite [Double
probability]
    if Double
probability Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
probability Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"dropout probability must be in [0,1)")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let draw :: ([Double], RNG) -> p -> ([Double], RNG)
draw ([Double]
xs, RNG
s) p
_ = let (Double
u, RNG
t) = RNG -> (Double, RNG)
uniform RNG
s in ([Double]
xs [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
probability then Double
0 else Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
probability)], RNG
t)
        ([Double]
mask, RNG
state) =
            if Training
mode Training -> Training -> Bool
forall a. Eq a => a -> a -> Bool
== Training
Inference
                then (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Tensor Double -> Int
forall a. Tensor a -> Int
T.numel Tensor Double
input) Double
1, RNG
rng)
                else (([Double], RNG) -> Int -> ([Double], RNG))
-> ([Double], RNG) -> [Int] -> ([Double], RNG)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' ([Double], RNG) -> Int -> ([Double], RNG)
forall {p}. ([Double], RNG) -> p -> ([Double], RNG)
draw ([], RNG
rng) [Int
1 .. Tensor Double -> Int
forall a. Tensor a -> Int
T.numel Tensor Double
input]
    m <- [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input) [Double]
mask
    y <- checkedMake (T.tensorShape input) (zipWith (*) (T.tensorValues input) mask)
    Right (y, m, state)

-- | Sample one inverted dropout mask per N,C pair and share it across every spatial position.
spatialDropout ::
    Training ->
    Double ->
    RNG ->
    T.Tensor Double ->
    Either MathError (T.Tensor Double, T.Tensor Double, RNG)
spatialDropout :: Training
-> Double
-> RNG
-> Tensor Double
-> Either MathError (Tensor Double, Tensor Double, RNG)
spatialDropout Training
mode Double
probability RNG
rng Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
n : Int
c : [Int]
spatial | Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
spatial) -> do
        prototype <- [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int
n, Int
c] (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c) Double
1)
        (_, mask, state) <- dropout mode probability rng prototype
        expanded <-
            T.tensor (n : c : replicate (length spatial) 1) (T.tensorValues mask)
                >>= T.broadcastTo (T.tensorShape input)
        output <-
            checkedMake (T.tensorShape input) (zipWith (*) (T.tensorValues input) (T.tensorValues expanded))
        Right (output, expanded, state)
    [Int]
_ -> MathError -> Either MathError (Tensor Double, Tensor Double, RNG)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"spatial dropout requires N,C,spatial")

{- | Per-channel running mean/variance and incorporated training-batch count for batch
normalization.
-}
data BatchNormState = BatchNormState
    { BatchNormState -> [Double]
runningMean :: [Double]
    -- ^ Per-channel running mean used during inference.
    , BatchNormState -> [Double]
runningVariance :: [Double]
    -- ^ Per-channel running variance used during inference.
    , BatchNormState -> Int
batchesSeen :: Int
    -- ^ Number of training batches incorporated into running statistics.
    }
    deriving (BatchNormState -> BatchNormState -> Bool
(BatchNormState -> BatchNormState -> Bool)
-> (BatchNormState -> BatchNormState -> Bool) -> Eq BatchNormState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BatchNormState -> BatchNormState -> Bool
== :: BatchNormState -> BatchNormState -> Bool
$c/= :: BatchNormState -> BatchNormState -> Bool
/= :: BatchNormState -> BatchNormState -> Bool
Eq, Int -> BatchNormState -> ShowS
[BatchNormState] -> ShowS
BatchNormState -> String
(Int -> BatchNormState -> ShowS)
-> (BatchNormState -> String)
-> ([BatchNormState] -> ShowS)
-> Show BatchNormState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BatchNormState -> ShowS
showsPrec :: Int -> BatchNormState -> ShowS
$cshow :: BatchNormState -> String
show :: BatchNormState -> String
$cshowList :: [BatchNormState] -> ShowS
showList :: [BatchNormState] -> ShowS
Show)

{- | Population variance normalizes a training batch; unbiased sample variance
updates runningVariance (a one-element batch uses population variance).
-}
batchNorm ::
    Training ->
    Double ->
    Double ->
    [Double] ->
    [Double] ->
    BatchNormState ->
    T.Tensor Double ->
    Either MathError (T.Tensor Double, BatchNormState)
batchNorm :: Training
-> Double
-> Double
-> [Double]
-> [Double]
-> BatchNormState
-> Tensor Double
-> Either MathError (Tensor Double, BatchNormState)
batchNorm Training
mode Double
momentum Double
eps [Double]
gamma [Double]
beta BatchNormState
state Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
n : Int
c : [Int]
spatial -> do
        Tensor Double -> Either MathError ()
validTensor Tensor Double
input
        [Double] -> Either MathError ()
finite (Double
momentum Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Double
eps Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
gamma [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
beta [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ BatchNormState -> [Double]
runningMean BatchNormState
state [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ BatchNormState -> [Double]
runningVariance BatchNormState
state)
        if Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
            Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
spatial Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
            Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
gamma Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c
            Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
beta Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c
            Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (BatchNormState -> [Double]
runningMean BatchNormState
state) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c
            Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (BatchNormState -> [Double]
runningVariance BatchNormState
state) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c
            Bool -> Bool -> Bool
|| Double
momentum Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
            Bool -> Bool -> Bool
|| Double
momentum Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1
            Bool -> Bool -> Bool
|| Double
eps Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
            Bool -> Bool -> Bool
|| BatchNormState -> Int
batchesSeen BatchNormState
state Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0) (BatchNormState -> [Double]
runningVariance BatchNormState
state)
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"batch normalization shape/state/options")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let perChannel :: Int
perChannel = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
spatial
            groups :: [[Double]]
groups =
                [ [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input (Int
batch Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
channel Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ix) | Int
batch <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
ix <- [Int] -> [[Int]]
indices [Int]
spatial]
                | Int
channel <- [Int
0 .. Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
            means :: [Double]
means = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\[Double]
xs -> [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
perChannel) [[Double]]
groups
            variances :: [Double]
variances = (Double -> [Double] -> Double)
-> [Double] -> [[Double]] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
m [Double]
xs -> [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
x <- [Double]
xs] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
perChannel) [Double]
means [[Double]]
groups
            ([Double]
usedMean, [Double]
usedVariance) = if Training
mode Training -> Training -> Bool
forall a. Eq a => a -> a -> Bool
== Training
Training then ([Double]
means, [Double]
variances) else (BatchNormState -> [Double]
runningMean BatchNormState
state, BatchNormState -> [Double]
runningVariance BatchNormState
state)
            update :: Double -> Double -> Double
update Double
old Double
new = (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
momentum) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
old Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
momentum Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
new
            unbiased :: [Double]
unbiased =
                (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map
                    (Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Int
perChannel Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 then Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
perChannel Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
perChannel Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) else Double
1))
                    [Double]
variances
            state' :: BatchNormState
state' =
                if Training
mode Training -> Training -> Bool
forall a. Eq a => a -> a -> Bool
== Training
Training
                    then
                        [Double] -> [Double] -> Int -> BatchNormState
BatchNormState
                            ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
update (BatchNormState -> [Double]
runningMean BatchNormState
state) [Double]
means)
                            ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
update (BatchNormState -> [Double]
runningVariance BatchNormState
state) [Double]
unbiased)
                            (BatchNormState -> Int
batchesSeen BatchNormState
state Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    else BatchNormState
state
        output <-
            [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake
                (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input)
                [ ([Double]
gamma [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int]
ix Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double]
usedMean [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt ([Double]
usedVariance [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)
                    Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
beta [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel
                | ix :: [Int]
ix@(Int
_ : Int
channel : [Int]
_) <- [Int] -> [[Int]]
indices (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input)
                ]
        Right (output, state')
    [Int]
_ -> MathError -> Either MathError (Tensor Double, BatchNormState)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"batch normalization requires N,C,...")

normalizeBy ::
    Bool -> Double -> ([Int] -> Int) -> T.Tensor Double -> Either MathError (T.Tensor Double)
normalizeBy :: Bool
-> Double
-> ([Int] -> Int)
-> Tensor Double
-> Either MathError (Tensor Double)
normalizeBy Bool
rms Double
eps [Int] -> Int
group Tensor Double
input = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    [Double] -> Either MathError ()
finite [Double
eps]
    if Double
eps Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Tensor Double -> Int
forall a. Tensor a -> Int
T.numel Tensor Double
input Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"normalization epsilon/empty input")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let allIx :: [[Int]]
allIx = [Int] -> [[Int]]
indices (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input)
        ids :: [Int]
ids = ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
group [[Int]]
allIx
        stats :: [(Int, (Double, Double))]
stats =
            [ let xs :: [Double]
xs = [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int]
ix | [Int]
ix <- [[Int]]
allIx, [Int] -> Int
group [Int]
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
g]
                  mean :: Double
mean = if Bool
rms then Double
0 else [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
                  variance :: Double
variance = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mean) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
x <- [Double]
xs] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
               in (Int
g, (Double
mean, Double
variance))
            | Int
g <- [Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
ids
            ]
        values :: [Double]
values =
            [ let (Double
mean, Double
variance) = (Double, Double)
-> ((Double, Double) -> (Double, Double))
-> Maybe (Double, Double)
-> (Double, Double)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Double
0, Double
0) (Double, Double) -> (Double, Double)
forall a. a -> a
id (Int -> [(Int, (Double, Double))] -> Maybe (Double, Double)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ([Int] -> Int
group [Int]
ix) [(Int, (Double, Double))]
stats)
               in (Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int]
ix Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mean) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
variance Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)
            | [Int]
ix <- [[Int]]
allIx
            ]
    [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input) [Double]
values

-- | Normalize the selected number of trailing axes independently for each remaining prefix index.
layerNorm :: Int -> Double -> T.Tensor Double -> Either MathError (T.Tensor Double)
layerNorm :: Int -> Double -> Tensor Double -> Either MathError (Tensor Double)
layerNorm Int
axes Double
eps Tensor Double
input
    | Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape = MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"layer norm trailing axis count")
    | Bool
otherwise =
        Bool
-> Double
-> ([Int] -> Int)
-> Tensor Double
-> Either MathError (Tensor Double)
normalizeBy Bool
False Double
eps ([Int] -> [Int] -> Int
flat (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
axes) [Int]
shape) ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
axes)) Tensor Double
input
  where
    shape :: [Int]
shape = Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input

-- | Divide trailing-axis groups by sqrt(mean(x*x)+epsilon), without subtracting a mean.
rmsNorm :: Int -> Double -> T.Tensor Double -> Either MathError (T.Tensor Double)
rmsNorm :: Int -> Double -> Tensor Double -> Either MathError (Tensor Double)
rmsNorm Int
axes Double
eps Tensor Double
input
    | Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape = MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"RMS norm trailing axis count")
    | Bool
otherwise =
        Bool
-> Double
-> ([Int] -> Int)
-> Tensor Double
-> Either MathError (Tensor Double)
normalizeBy Bool
True Double
eps ([Int] -> [Int] -> Int
flat (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
axes) [Int]
shape) ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
axes)) Tensor Double
input
  where
    shape :: [Int]
shape = Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input

-- | Normalize spatial axes separately per batch/channel in N,C,spatial layout.
instanceNorm :: Double -> T.Tensor Double -> Either MathError (T.Tensor Double)
instanceNorm :: Double -> Tensor Double -> Either MathError (Tensor Double)
instanceNorm Double
eps Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
_ : Int
c : [Int]
spatial | Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
spatial) -> Bool
-> Double
-> ([Int] -> Int)
-> Tensor Double
-> Either MathError (Tensor Double)
normalizeBy Bool
False Double
eps (\[Int]
ix -> [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) Tensor Double
input
    [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"instance norm requires N,C,spatial")

-- | Normalize channels/spatial positions per batch and channel group; channel count must be divisible by groups.
groupNorm :: Int -> Double -> T.Tensor Double -> Either MathError (T.Tensor Double)
groupNorm :: Int -> Double -> Tensor Double -> Either MathError (Tensor Double)
groupNorm Int
groups Double
eps Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
_ : Int
c : [Int]
_
        | Int
groups Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
c Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
groups Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 ->
            Bool
-> Double
-> ([Int] -> Int)
-> Tensor Double
-> Either MathError (Tensor Double)
normalizeBy Bool
False Double
eps (\[Int]
ix -> [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
groups Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
c Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups)) Tensor Double
input
    [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"group norm channel/group shape")

-- | Normalize each selected-axis slice and multiply it by the matching explicit magnitude; zero directions fail.
weightNorm :: Int -> [Double] -> T.Tensor Double -> Either MathError (T.Tensor Double)
weightNorm :: Int
-> [Double] -> Tensor Double -> Either MathError (Tensor Double)
weightNorm Int
axis [Double]
magnitudes Tensor Double
input = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    [Double] -> Either MathError ()
finite [Double]
magnitudes
    let shape :: [Int]
shape = Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input
    if Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
magnitudes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int]
shape [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"weight norm axis/magnitude shape")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let allIx :: [[Int]]
allIx = [Int] -> [[Int]]
indices [Int]
shape
        norms :: [Double]
norms =
            [ Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int]
ix Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | [Int]
ix <- [[Int]]
allIx, [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i])
            | Int
i <- [Int
0 .. [Int]
shape [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]
    if (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0) [Double]
norms
        then MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"weight normalization zero direction")
        else
            [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake
                [Int]
shape
                [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int]
ix Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
magnitudes [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double]
norms [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)) | [Int]
ix <- [[Int]]
allIx]

{- | Weight initializer selector: Xavier/Kaiming uniform/normal,Orthogonal gain,or
TruncatedNormal mean sigma cutoff. initializeWeights also takes explicit fan counts and
RNG.
-}
data Initializer
    = XavierUniform
    | XavierNormal
    | KaimingUniform
    | KaimingNormal
    | Orthogonal Double
    | TruncatedNormal Double Double Double
    deriving (Initializer -> Initializer -> Bool
(Initializer -> Initializer -> Bool)
-> (Initializer -> Initializer -> Bool) -> Eq Initializer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Initializer -> Initializer -> Bool
== :: Initializer -> Initializer -> Bool
$c/= :: Initializer -> Initializer -> Bool
/= :: Initializer -> Initializer -> Bool
Eq, Int -> Initializer -> ShowS
[Initializer] -> ShowS
Initializer -> String
(Int -> Initializer -> ShowS)
-> (Initializer -> String)
-> ([Initializer] -> ShowS)
-> Show Initializer
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Initializer -> ShowS
showsPrec :: Int -> Initializer -> ShowS
$cshow :: Initializer -> String
show :: Initializer -> String
$cshowList :: [Initializer] -> ShowS
showList :: [Initializer] -> ShowS
Show)

-- | Generate an immutable tensor and advanced RNG. Fans are explicit; orthogonal initialization requires a positive matrix shape.
initializeWeights ::
    Initializer -> Int -> Int -> [Int] -> RNG -> Either MathError (T.Tensor Double, RNG)
initializeWeights :: Initializer
-> Int
-> Int
-> [Int]
-> RNG
-> Either MathError (Tensor Double, RNG)
initializeWeights Initializer
kind Int
fanIn Int
fanOut [Int]
shape RNG
rng = do
    count <- [Int] -> Either MathError Int
T.shapeSize [Int]
shape
    if fanIn < 1 || fanOut < 1
        then Left (InvalidArgument "positive initialization fan counts")
        else Right ()
    let std = case Initializer
kind of
            Initializer
XavierNormal -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanIn Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanOut))
            Initializer
_ -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanIn)
        bound = case Initializer
kind of
            Initializer
XavierUniform -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
6 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanIn Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanOut))
            Initializer
_ -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
6 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
fanIn)
        uniformDraw RNG
s = let (Double
u, RNG
t) = RNG -> (Double, RNG)
uniform RNG
s in (Double, RNG) -> Either a (Double, RNG)
forall a b. b -> Either a b
Right ((Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
bound, RNG
t)
    (values, state) <- case kind of
        Initializer
XavierUniform -> Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count RNG -> Either MathError (Double, RNG)
forall {a}. RNG -> Either a (Double, RNG)
uniformDraw RNG
rng
        Initializer
KaimingUniform -> Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count RNG -> Either MathError (Double, RNG)
forall {a}. RNG -> Either a (Double, RNG)
uniformDraw RNG
rng
        Initializer
XavierNormal -> Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count (Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
0 Double
std) RNG
rng
        Initializer
KaimingNormal -> Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count (Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
0 Double
std) RNG
rng
        TruncatedNormal Double
mean Double
sigma Double
limit -> do
            [Double] -> Either MathError ()
finite [Double
mean, Double
sigma, Double
limit]
            if Double
sigma Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
limit Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
                then MathError -> Either MathError ([Double], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"truncated normal sigma/limit")
                else Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count (Double
-> Double -> Double -> Int -> RNG -> Either MathError (Double, RNG)
rejection Double
mean Double
sigma Double
limit Int
0) RNG
rng
        Orthogonal Double
gain -> do
            [Double] -> Either MathError ()
finite [Double
gain]
            case [Int]
shape of
                [Int
rows, Int
cols] | Int
rows Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
cols Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
                    (draws, s) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
count (Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
0 Double
1) RNG
rng
                    if rows >= cols
                        then do
                            qs <- orthogonalize [[draws !! (i * cols + j) | i <- [0 .. rows - 1]] | j <- [0 .. cols - 1]]
                            Right ([gain * ((qs !! j) !! i) | i <- [0 .. rows - 1], j <- [0 .. cols - 1]], s)
                        else do
                            qs <- orthogonalize [[draws !! (i * cols + j) | j <- [0 .. cols - 1]] | i <- [0 .. rows - 1]]
                            Right (map (* gain) (concat qs), s)
                [Int]
_ -> MathError -> Either MathError ([Double], RNG)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"orthogonal initialization requires positive matrix shape")
    output <- checkedMake shape values
    Right (output, state)
  where
    rejection :: Double
-> Double -> Double -> Int -> RNG -> Either MathError (Double, RNG)
rejection Double
mean Double
sigma Double
limit Int
attempts RNG
s
        | Int
attempts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
100000 = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"truncated normal rejection" Int
attempts)
        | Bool
otherwise = do
            (x, t) <- Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
mean Double
sigma RNG
s
            if abs (x - mean) <= limit * sigma
                then Right (x, t)
                else rejection mean sigma limit (attempts + 1) t
    orthogonalize :: [[Double]] -> Either MathError [[Double]]
orthogonalize =
        ([[Double]] -> [Double] -> Either MathError [[Double]])
-> [[Double]] -> [[Double]] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
            ( \[[Double]]
basis [Double]
v -> do
                let residual :: [Double]
residual = ([Double] -> [Double] -> [Double])
-> [Double] -> [[Double]] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\[Double]
r [Double]
q -> (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Double]
r ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((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]
r [Double]
q))) [Double]
q)) [Double]
v [[Double]]
basis
                    norm :: Double
norm = Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x) [Double]
residual))
                if Double
norm Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-14
                    then MathError -> Either MathError [[Double]]
forall a b. a -> Either a b
Left (Int -> MathError
RankDeficient ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
basis))
                    else [[Double]] -> Either MathError [[Double]]
forall a b. b -> Either a b
Right ([[Double]]
basis [[Double]] -> [[Double]] -> [[Double]]
forall a. [a] -> [a] -> [a]
++ [(Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
norm) [Double]
residual])
            )
            []

-- | Gather checked row indices from [vocabulary,features] weights in requested order.
embedding :: T.Tensor Double -> [Int] -> Either MathError (T.Tensor Double)
embedding :: Tensor Double -> [Int] -> Either MathError (Tensor Double)
embedding Tensor Double
weights [Int]
ids = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
weights
    case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
weights of
        [Int
rows, Int
cols]
            | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rows) [Int]
ids ->
                [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake [[Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids, Int
cols] [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
weights [Int
i, Int
j] | Int
i <- [Int]
ids, Int
j <- [Int
0 .. Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"embedding weights/indices")

denseGeneric :: (Num a) => T.Tensor a -> T.Tensor a -> [a] -> Either MathError (T.Tensor a)
denseGeneric :: forall a.
Num a =>
Tensor a -> Tensor a -> [a] -> Either MathError (Tensor a)
denseGeneric Tensor a
input Tensor a
weights [a]
bias = case (Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
input, Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
weights) of
    ([Int
n, Int
features], [Int
out, Int
inner])
        | Int
features Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
inner Bool -> Bool -> Bool
&& [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
bias Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
out ->
            [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor
                [Int
n, Int
out]
                [ [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
get Tensor a
input [Int
i, Int
k] a -> a -> a
forall a. Num a => a -> a -> a
* Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
get Tensor a
weights [Int
j, Int
k] | Int
k <- [Int
0 .. Int
features Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
bias [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                , Int
j <- [Int
0 .. Int
out Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
    ([Int], [Int])
_ -> MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"dense expects [batch,in],[out,in],[out]")

-- | Apply x*transpose(weights)+bias to [batch,in] inputs and [out,in] weights.
dense :: T.Tensor Double -> T.Tensor Double -> [Double] -> Either MathError (T.Tensor Double)
dense :: Tensor Double
-> Tensor Double -> [Double] -> Either MathError (Tensor Double)
dense Tensor Double
input Tensor Double
weights [Double]
bias = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    Tensor Double -> Either MathError ()
validTensor Tensor Double
weights
    [Double] -> Either MathError ()
finite [Double]
bias
    y <- Tensor Double
-> Tensor Double -> [Double] -> Either MathError (Tensor Double)
forall a.
Num a =>
Tensor a -> Tensor a -> [a] -> Either MathError (Tensor a)
denseGeneric Tensor Double
input Tensor Double
weights [Double]
bias
    validTensor y
    Right y

-- | Symbolic dense layer with gradients through inputs, weights, and explicit bias expressions.
denseAD :: AD.ADTensor -> AD.ADTensor -> [AD.Expr] -> Either MathError AD.ADTensor
denseAD :: ADTensor -> ADTensor -> [Expr] -> Either MathError ADTensor
denseAD ADTensor
input ADTensor
weights [Expr]
bias = do
    x <- [Int] -> [Expr] -> Either MathError (Tensor Expr)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor (ADTensor -> [Int]
AD.adShape ADTensor
input) (ADTensor -> [Expr]
AD.adValues ADTensor
input)
    w <- T.tensor (AD.adShape weights) (AD.adValues weights)
    y <- denseGeneric x w bias
    AD.adTensor (T.tensorShape y) (T.tensorValues y)

-- | Per-axis stride,padding,dilation,and channel-group settings for checked convolution.
data ConvOptions = ConvOptions
    { ConvOptions -> [Int]
convolutionStride :: [Int]
    -- ^ Positive spatial stride per convolution axis.
    , ConvOptions -> [Int]
convolutionPadding :: [Int]
    -- ^ Nonnegative symmetric padding per spatial axis.
    , ConvOptions -> [Int]
convolutionDilation :: [Int]
    -- ^ Positive spacing between kernel elements per spatial axis.
    , ConvOptions -> Int
convolutionGroups :: Int
    -- ^ Positive group count dividing input/output channels for grouped convolution.
    }
    deriving (ConvOptions -> ConvOptions -> Bool
(ConvOptions -> ConvOptions -> Bool)
-> (ConvOptions -> ConvOptions -> Bool) -> Eq ConvOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ConvOptions -> ConvOptions -> Bool
== :: ConvOptions -> ConvOptions -> Bool
$c/= :: ConvOptions -> ConvOptions -> Bool
/= :: ConvOptions -> ConvOptions -> Bool
Eq, Int -> ConvOptions -> ShowS
[ConvOptions] -> ShowS
ConvOptions -> String
(Int -> ConvOptions -> ShowS)
-> (ConvOptions -> String)
-> ([ConvOptions] -> ShowS)
-> Show ConvOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConvOptions -> ShowS
showsPrec :: Int -> ConvOptions -> ShowS
$cshow :: ConvOptions -> String
show :: ConvOptions -> String
$cshowList :: [ConvOptions] -> ShowS
showList :: [ConvOptions] -> ShowS
Show)

convGeneric ::
    (Num a) =>
    Bool -> [Int] -> ConvOptions -> T.Tensor a -> T.Tensor a -> [a] -> Either MathError (T.Tensor a)
convGeneric :: forall a.
Num a =>
Bool
-> [Int]
-> ConvOptions
-> Tensor a
-> Tensor a
-> [a]
-> Either MathError (Tensor a)
convGeneric Bool
transposed [Int]
outputPadding ConvOptions
options Tensor a
input Tensor a
weights [a]
bias = case (Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
input, Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
weights) of
    (Int
n : Int
cin : [Int]
spatial, Int
first : Int
second : [Int]
kernel) -> do
        let dims :: Int
dims = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
spatial
            strides :: [Int]
strides = ConvOptions -> [Int]
convolutionStride ConvOptions
options
            padding :: [Int]
padding = ConvOptions -> [Int]
convolutionPadding ConvOptions
options
            dilation :: [Int]
dilation = ConvOptions -> [Int]
convolutionDilation ConvOptions
options
            groups :: Int
groups = ConvOptions -> Int
convolutionGroups ConvOptions
options
            coutInteger :: Integer
coutInteger = if Bool
transposed then Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
second Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
groups else Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
first
            cout :: Int
cout = Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
coutInteger
        if Integer
coutInteger Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
            Bool -> Bool -> Bool
|| Integer
coutInteger Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1
            Bool -> Bool -> Bool
|| Int
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
            Bool -> Bool -> Bool
|| Int
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
3
            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
kernel Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims
            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
strides Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims
            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
padding Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims
            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dilation Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims
            Bool -> Bool -> Bool
|| Int
groups Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
            Bool -> Bool -> Bool
|| Int
cin Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
            Bool -> Bool -> Bool
|| Int
cout Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
            Bool -> Bool -> Bool
|| Int
cin Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
groups Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
            Bool -> Bool -> Bool
|| Int
cout Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
groups Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
            Bool -> Bool -> Bool
|| [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
bias Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
cout
            Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) ([Int]
kernel [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
strides [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
dilation [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
spatial)
            Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
padding
            Bool -> Bool -> Bool
|| ( if Bool
transposed
                    then
                        Int
first Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
cin
                            Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
outputPadding Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims
                            Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
p Int
s -> Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
s) [Int]
outputPadding [Int]
strides)
                    else Int
second Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
cin Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups
               )
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"convolution layout, channels, groups, stride, or padding")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let integerSpatial :: [Integer]
integerSpatial = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
spatial
            integerKernel :: [Integer]
integerKernel = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
kernel
            integerStride :: [Integer]
integerStride = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
strides
            integerPadding :: [Integer]
integerPadding = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
padding
            integerDilation :: [Integer]
integerDilation = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
dilation
            outputDimensions :: [Integer]
outputDimensions =
                if Bool
transposed
                    then
                        (Integer -> Integer -> Integer -> Integer -> Integer -> Integer)
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
forall {t} {t} {t} {t} {t} {a}.
(t -> t -> t -> t -> t -> a)
-> [t] -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith5
                            (\Integer
s Integer
k Integer
st Integer
p Integer
d -> (Integer
s Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
st Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
                            [Integer]
integerSpatial
                            [Integer]
integerKernel
                            [Integer]
integerStride
                            [Integer]
integerPadding
                            [Integer]
integerDilation
                            [Integer] -> [Integer] -> [Integer]
`zipAdd` (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
outputPadding
                    else
                        (Integer -> Integer -> Integer -> Integer -> Integer -> Integer)
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
-> [Integer]
forall {t} {t} {t} {t} {t} {a}.
(t -> t -> t -> t -> t -> a)
-> [t] -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith5
                            (\Integer
s Integer
k Integer
st Integer
p Integer
d -> (Integer
s Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
st Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
                            [Integer]
integerSpatial
                            [Integer]
integerKernel
                            [Integer]
integerStride
                            [Integer]
integerPadding
                            [Integer]
integerDilation
        outShape <- [Integer] -> Either MathError [Int]
checkedDimensions [Integer]
outputDimensions
        if any (<= 0) outShape then Left (InvalidArgument "convolution empty output") else Right ()
        _ <- T.shapeSize (n : cout : outShape)
        let values =
                [ let group :: Int
group = Int
oc Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
cout Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups)
                      contributions :: [a]
contributions =
                        [ let (Bool
valid, [Int]
ip) =
                                if Bool
transposed
                                    then
                                        let numerator :: [Integer]
numerator =
                                                (Integer -> Integer -> Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer] -> [Integer] -> [Integer]
forall {t} {t} {t} {t} {a}.
(t -> t -> t -> t -> a) -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith4
                                                    (\Integer
o Integer
p Integer
k Integer
d -> Integer
o Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d)
                                                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
op)
                                                    [Integer]
integerPadding
                                                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
kp)
                                                    [Integer]
integerDilation
                                            point :: [Integer]
point = (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div [Integer]
numerator [Integer]
integerStride
                                         in ( [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Integer -> Integer -> Bool) -> [Integer] -> [Integer] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Integer
v Integer
st -> Integer
v Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
st Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0) [Integer]
numerator [Integer]
integerStride) Bool -> Bool -> Bool
&& [Integer] -> [Integer] -> Bool
forall {a}. (Ord a, Num a) => [a] -> [a] -> Bool
inside [Integer]
point [Integer]
integerSpatial
                                            , (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
point
                                            )
                                    else
                                        let point :: [Integer]
point =
                                                (Integer -> Integer -> Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer] -> [Integer] -> [Integer]
forall {t} {t} {t} {t} {a}.
(t -> t -> t -> t -> a) -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith4
                                                    (\Integer
o Integer
st Integer
k Integer
d -> Integer
o Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
st Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d)
                                                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
op)
                                                    [Integer]
integerStride
                                                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
kp)
                                                    [Integer]
integerDilation
                                                    [Integer] -> [Integer] -> [Integer]
`zipSubtract` [Integer]
integerPadding
                                         in ([Integer] -> [Integer] -> Bool
forall {a}. (Ord a, Num a) => [a] -> [a] -> Bool
inside [Integer]
point [Integer]
integerSpatial, (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
point)
                              channel :: Int
channel = Int
group Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
cin Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ic
                              weightIndex :: [Int]
weightIndex = if Bool
transposed then Int
channel Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Int
oc Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` (Int
cout Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups)) Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
kp else Int
oc Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
ic Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
kp
                           in if Bool
valid then Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
get Tensor a
input (Int
batch Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
channel Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ip) a -> a -> a
forall a. Num a => a -> a -> a
* Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
get Tensor a
weights [Int]
weightIndex else a
0
                        | Int
ic <- [Int
0 .. Int
cin Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
groups Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        , [Int]
kp <- [Int] -> [[Int]]
indices [Int]
kernel
                        ]
                   in [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
contributions a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
bias [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
oc
                | Int
batch <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                , Int
oc <- [Int
0 .. Int
cout Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                , [Int]
op <- [Int] -> [[Int]]
indices [Int]
outShape
                ]
        T.tensor (n : cout : outShape) values
    ([Int], [Int])
_ -> MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"convolution requires N,C,spatial and two channel weight axes")
  where
    inside :: [a] -> [a] -> Bool
inside [a]
point [a]
shape = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((a -> a -> Bool) -> [a] -> [a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
i a
d -> a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
d) [a]
point [a]
shape)
    zipAdd :: [Integer] -> [Integer] -> [Integer]
zipAdd = (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+)
    zipSubtract :: [Integer] -> [Integer] -> [Integer]
zipSubtract = (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-)
    zipWith4 :: (t -> t -> t -> t -> a) -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith4 t -> t -> t -> t -> a
f [t]
as [t]
bs [t]
cs [t]
ds = [t -> t -> t -> t -> a
f t
a t
b t
c t
d | (t
a, (t
b, (t
c, t
d))) <- [t] -> [(t, (t, t))] -> [(t, (t, (t, t)))]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
as ([t] -> [(t, t)] -> [(t, (t, t))]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
bs ([t] -> [t] -> [(t, t)]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
cs [t]
ds))]
    zipWith5 :: (t -> t -> t -> t -> t -> a)
-> [t] -> [t] -> [t] -> [t] -> [t] -> [a]
zipWith5 t -> t -> t -> t -> t -> a
f [t]
as [t]
bs [t]
cs [t]
ds [t]
es = [t -> t -> t -> t -> t -> a
f t
a t
b t
c t
d t
e | (t
a, (t
b, (t
c, (t
d, t
e)))) <- [t] -> [(t, (t, (t, t)))] -> [(t, (t, (t, (t, t))))]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
as ([t] -> [(t, (t, t))] -> [(t, (t, (t, t)))]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
bs ([t] -> [(t, t)] -> [(t, (t, t))]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
cs ([t] -> [t] -> [(t, t)]
forall a b. [a] -> [b] -> [(a, b)]
zip [t]
ds [t]
es)))]

-- | Direct checked 1D/2D/3D cross-correlation with grouped channels and explicit stride/padding/dilation.
convolution ::
    ConvOptions -> T.Tensor Double -> T.Tensor Double -> [Double] -> Either MathError (T.Tensor Double)
convolution :: ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
convolution ConvOptions
options Tensor Double
input Tensor Double
weights [Double]
bias = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    Tensor Double -> Either MathError ()
validTensor Tensor Double
weights
    [Double] -> Either MathError ()
finite [Double]
bias
    output <- Bool
-> [Int]
-> ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
forall a.
Num a =>
Bool
-> [Int]
-> ConvOptions
-> Tensor a
-> Tensor a
-> [a]
-> Either MathError (Tensor a)
convGeneric Bool
False [] ConvOptions
options Tensor Double
input Tensor Double
weights [Double]
bias
    validTensor output
    Right output

-- | Direct adjoint-layout convolution; output padding is per-axis and strictly less than stride.
transposedConvolution ::
    [Int] ->
    ConvOptions ->
    T.Tensor Double ->
    T.Tensor Double ->
    [Double] ->
    Either MathError (T.Tensor Double)
transposedConvolution :: [Int]
-> ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
transposedConvolution [Int]
extra ConvOptions
options Tensor Double
input Tensor Double
weights [Double]
bias = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    Tensor Double -> Either MathError ()
validTensor Tensor Double
weights
    [Double] -> Either MathError ()
finite [Double]
bias
    output <- Bool
-> [Int]
-> ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
forall a.
Num a =>
Bool
-> [Int]
-> ConvOptions
-> Tensor a
-> Tensor a
-> [a]
-> Either MathError (Tensor a)
convGeneric Bool
True [Int]
extra ConvOptions
options Tensor Double
input Tensor Double
weights [Double]
bias
    validTensor output
    Right output

-- | Set groups to input channels; output channel multipliers remain explicit in the weight shape.
depthwiseConvolution ::
    ConvOptions -> T.Tensor Double -> T.Tensor Double -> [Double] -> Either MathError (T.Tensor Double)
depthwiseConvolution :: ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
depthwiseConvolution ConvOptions
options Tensor Double
input Tensor Double
weights [Double]
bias = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
_ : Int
channels : [Int]
_ -> ConvOptions
-> Tensor Double
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
convolution (ConvOptions
options{convolutionGroups = channels}) Tensor Double
input Tensor Double
weights [Double]
bias
    [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"depthwise input channel shape")

-- | Construct the same direct convolution from scalar graph arithmetic, with automatic input/kernel/bias gradients.
convolutionAD ::
    ConvOptions -> AD.ADTensor -> AD.ADTensor -> [AD.Expr] -> Either MathError AD.ADTensor
convolutionAD :: ConvOptions
-> ADTensor -> ADTensor -> [Expr] -> Either MathError ADTensor
convolutionAD ConvOptions
options ADTensor
input ADTensor
weights [Expr]
bias = do
    x <- [Int] -> [Expr] -> Either MathError (Tensor Expr)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor (ADTensor -> [Int]
AD.adShape ADTensor
input) (ADTensor -> [Expr]
AD.adValues ADTensor
input)
    w <- T.tensor (AD.adShape weights) (AD.adValues weights)
    output <- convGeneric False [] options x w bias
    AD.adTensor (T.tensorShape output) (T.tensorValues output)

-- | Pooling reduction selector: MaxPool or AveragePool over each checked window.
data PoolMode = MaxPool | AveragePool deriving (PoolMode -> PoolMode -> Bool
(PoolMode -> PoolMode -> Bool)
-> (PoolMode -> PoolMode -> Bool) -> Eq PoolMode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PoolMode -> PoolMode -> Bool
== :: PoolMode -> PoolMode -> Bool
$c/= :: PoolMode -> PoolMode -> Bool
/= :: PoolMode -> PoolMode -> Bool
Eq, Int -> PoolMode -> ShowS
[PoolMode] -> ShowS
PoolMode -> String
(Int -> PoolMode -> ShowS)
-> (PoolMode -> String) -> ([PoolMode] -> ShowS) -> Show PoolMode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PoolMode -> ShowS
showsPrec :: Int -> PoolMode -> ShowS
$cshow :: PoolMode -> String
show :: PoolMode -> String
$cshowList :: [PoolMode] -> ShowS
showList :: [PoolMode] -> ShowS
Show)

-- | Windowed max/average pooling. Padded elements are excluded from the average; entirely padded windows fail.
pooling ::
    PoolMode -> [Int] -> [Int] -> [Int] -> T.Tensor Double -> Either MathError (T.Tensor Double)
pooling :: PoolMode
-> [Int]
-> [Int]
-> [Int]
-> Tensor Double
-> Either MathError (Tensor Double)
pooling PoolMode
mode [Int]
kernel [Int]
strides [Int]
padding Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
n : Int
c : [Int]
spatial -> do
        Tensor Double -> Either MathError ()
validTensor Tensor Double
input
        let dims :: Int
dims = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
spatial
        if Int
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
            Bool -> Bool -> Bool
|| Int
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
3
            Bool -> Bool -> Bool
|| ([Int] -> Bool) -> [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dims) (Int -> Bool) -> ([Int] -> Int) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Int]
kernel, [Int]
strides, [Int]
padding]
            Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) ([Int]
kernel [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
strides [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
spatial)
            Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
padding
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"pooling shapes/options")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let dimensions :: [Integer]
dimensions =
                (Integer -> Integer -> (Integer, Integer) -> Integer)
-> [Integer] -> [Integer] -> [(Integer, Integer)] -> [Integer]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
                    (\Integer
s Integer
k (Integer
st, Integer
p) -> (Integer
s Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
st Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
spatial)
                    ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
kernel)
                    ([Integer] -> [Integer] -> [(Integer, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
strides) ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
padding))
        outShape <- [Integer] -> Either MathError [Int]
checkedDimensions [Integer]
dimensions
        if any (<= 0) outShape then Left (InvalidArgument "pooling empty output") else Right ()
        _ <- T.shapeSize (n : c : outShape)
        vals <-
            mapM
                ( \(Int
batch, Int
channel, [Int]
op) -> do
                    let samples :: [Double]
samples =
                            [ Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input (Int
batch Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
channel Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
ip)
                            | [Int]
kp <- [Int] -> [[Int]]
indices [Int]
kernel
                            , let ip :: [Integer]
ip =
                                    (Integer -> (Integer, Integer) -> Integer -> Integer)
-> [Integer] -> [(Integer, Integer)] -> [Integer] -> [Integer]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
                                        (\Integer
o (Integer
st, Integer
p) Integer
k -> Integer
o Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
st Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
k)
                                        ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
op)
                                        ([Integer] -> [Integer] -> [(Integer, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
strides) ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
padding))
                                        ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
kp)
                            , [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Integer -> Int -> Bool) -> [Integer] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Integer
i Int
d -> Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 Bool -> Bool -> Bool
&& Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d) [Integer]
ip [Int]
spatial)
                            ]
                    if [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
samples
                        then MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"pool window entirely padding")
                        else
                            Double -> Either MathError Double
forall a b. b -> Either a b
Right (if PoolMode
mode PoolMode -> PoolMode -> Bool
forall a. Eq a => a -> a -> Bool
== PoolMode
MaxPool then [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
samples else [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
samples Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
samples))
                )
                [(batch, channel, op) | batch <- [0 .. n - 1], channel <- [0 .. c - 1], op <- indices outShape]
        checkedMake (n : c : outShape) vals
    [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"pooling requires N,C,spatial")

-- | Partition spatial axes by floor/ceil output-bin boundaries, then reduce each bin.
adaptivePooling :: PoolMode -> [Int] -> T.Tensor Double -> Either MathError (T.Tensor Double)
adaptivePooling :: PoolMode
-> [Int] -> Tensor Double -> Either MathError (Tensor Double)
adaptivePooling PoolMode
mode [Int]
outShape Tensor Double
input = case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
    Int
n : Int
c : [Int]
spatial | Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
spatial) Bool -> Bool -> Bool
&& [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
outShape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
spatial Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) ([Int]
outShape [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
spatial) -> do
        Tensor Double -> Either MathError ()
validTensor Tensor Double
input
        _ <- [Int] -> Either MathError Int
T.shapeSize (Int
n Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
c Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
outShape)
        let value Int
batch Int
channel [a]
op =
                let points :: [[Int]]
points =
                        [[Int]] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence
                            [ let start :: Integer
start = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
o Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
s Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d
                                  end :: Integer
end = ((a -> Integer
forall a. Integral a => a -> Integer
toInteger a
o Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
s Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d
                               in [Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
start .. Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                            | (a
o, (Int
s, Int
d)) <- [a] -> [(Int, Int)] -> [(a, (Int, Int))]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
op ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
spatial [Int]
outShape)
                            ]
                    xs :: [Double]
xs = [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input (Int
batch Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
channel Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ip) | [Int]
ip <- [[Int]]
points]
                 in if PoolMode
mode PoolMode -> PoolMode -> Bool
forall a. Eq a => a -> a -> Bool
== PoolMode
MaxPool then [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs else [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
        checkedMake
            (n : c : outShape)
            [value batch channel op | batch <- [0 .. n - 1], channel <- [0 .. c - 1], op <- indices outShape]
    [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"adaptive pooling spatial shape")

-- | Reduce all spatial positions to one value per N,C pair.
globalPooling :: PoolMode -> T.Tensor Double -> Either MathError (T.Tensor Double)
globalPooling :: PoolMode -> Tensor Double -> Either MathError (Tensor Double)
globalPooling PoolMode
mode Tensor Double
input = PoolMode
-> [Int] -> Tensor Double -> Either MathError (Tensor Double)
adaptivePooling PoolMode
mode (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int
1) Tensor Double
input

-- | Return attended values and normalized weights for Q[N,D],K[M,D],V[M,F]; True masks permit keys.
scaledDotProductAttention ::
    T.Tensor Double ->
    T.Tensor Double ->
    T.Tensor Double ->
    Maybe [[Bool]] ->
    Maybe [[Double]] ->
    Either MathError (T.Tensor Double, T.Tensor Double)
scaledDotProductAttention :: Tensor Double
-> Tensor Double
-> Tensor Double
-> Maybe [[Bool]]
-> Maybe [[Double]]
-> Either MathError (Tensor Double, Tensor Double)
scaledDotProductAttention Tensor Double
query Tensor Double
key Tensor Double
value Maybe [[Bool]]
mask Maybe [[Double]]
bias = do
    (Tensor Double -> Either MathError ())
-> [Tensor Double] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Tensor Double -> Either MathError ()
validTensor [Tensor Double
query, Tensor Double
key, Tensor Double
value]
    case (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
query, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
key, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
value) of
        ([Int
n, Int
d], [Int
m, Int
k], [Int
v, Int
features]) | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k Bool -> Bool -> Bool
&& Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
            let allowed :: [[Bool]]
allowed = [[Bool]] -> ([[Bool]] -> [[Bool]]) -> Maybe [[Bool]] -> [[Bool]]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> [Bool] -> [[Bool]]
forall a. Int -> a -> [a]
replicate Int
n (Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Int
m Bool
True)) [[Bool]] -> [[Bool]]
forall a. a -> a
id Maybe [[Bool]]
mask
                additive :: [[Double]]
additive = [[Double]]
-> ([[Double]] -> [[Double]]) -> Maybe [[Double]] -> [[Double]]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate Int
n (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
m Double
0)) [[Double]] -> [[Double]]
forall a. a -> a
id Maybe [[Double]]
bias
            if [[Bool]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Bool]]
allowed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n
                Bool -> Bool -> Bool
|| ([Bool] -> Bool) -> [[Bool]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m) (Int -> Bool) -> ([Bool] -> Int) -> [Bool] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Bool]]
allowed
                Bool -> Bool -> Bool
|| [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
additive Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n
                Bool -> Bool -> Bool
|| ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Double]]
additive
                then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"attention mask/bias shape")
                else [Double] -> Either MathError ()
finite ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
additive)
            rows <-
                (Int -> Either MathError [Double])
-> [Int] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                    ( \Int
i -> do
                        let active :: [Int]
active = [Int
j | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], ([[Bool]]
allowed [[Bool]] -> Int -> [Bool]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) [Bool] -> Int -> Bool
forall a. HasCallStack => [a] -> Int -> a
!! Int
j]
                            scores :: [Double]
scores =
                                [ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
query [Int
i, Int
t] Double -> Double -> Double
forall a. Num a => a -> a -> a
* Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
key [Int
j, Int
t] | Int
t <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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
d)
                                    Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ([[Double]]
additive [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                                | Int
j <- [Int]
active
                                ]
                        probabilities <- if [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
active then [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [] else [Double] -> Either MathError [Double]
softmax [Double]
scores
                        Right [maybe 0 id (lookup j (zip active probabilities)) | j <- [0 .. m - 1]]
                    )
                    [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            output <-
                checkedMake
                    [n, features]
                    [ sum [(rows !! i) !! j * get value [j, t] | j <- [0 .. m - 1]]
                    | i <- [0 .. n - 1]
                    , t <- [0 .. features - 1]
                    ]
            weights <- checkedMake [n, m] (concat rows)
            Right (output, weights)
        ([Int], [Int], [Int])
_ -> MathError -> Either MathError (Tensor Double, Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"attention requires Q[N,D],K[M,D],V[M,F]")

{- | Split feature dimension into equal heads, attend independently, concatenate.
Projection matrices are explicit dense layers composed by the caller.
-}
multiheadAttention ::
    Int ->
    T.Tensor Double ->
    T.Tensor Double ->
    T.Tensor Double ->
    Maybe [[Bool]] ->
    Either MathError (T.Tensor Double)
multiheadAttention :: Int
-> Tensor Double
-> Tensor Double
-> Tensor Double
-> Maybe [[Bool]]
-> Either MathError (Tensor Double)
multiheadAttention Int
heads Tensor Double
q Tensor Double
k Tensor Double
v Maybe [[Bool]]
mask = case (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
q, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
k, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
v) of
    ([Int
n, Int
d], [Int
m, Int
kd], [Int
vm, Int
vd]) | Int
heads Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
kd Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
vd Bool -> Bool -> Bool
&& Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
vm Bool -> Bool -> Bool
&& Int
d Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
heads Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> do
        let width :: Int
width = Int
d Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
heads
            select :: Tensor a -> Int -> Int -> Either MathError (Tensor a)
select Tensor a
input Int
rows Int
headIndex =
                [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor
                    [Int
rows, Int
width]
                    [Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
get Tensor a
input [Int
i, Int
headIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j] | Int
i <- [Int
0 .. Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- [Int
0 .. Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        outputs <-
            (Int -> Either MathError (Tensor Double))
-> [Int] -> Either MathError [Tensor Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                ( \Int
h -> do
                    qh <- Tensor Double -> Int -> Int -> Either MathError (Tensor Double)
forall {a}. Tensor a -> Int -> Int -> Either MathError (Tensor a)
select Tensor Double
q Int
n Int
h
                    kh <- select k m h
                    vh <- select v m h
                    fst <$> scaledDotProductAttention qh kh vh mask Nothing
                )
                [Int
0 .. Int
heads Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        checkedMake
            [n, d]
            [get (outputs !! h) [i, j] | i <- [0 .. n - 1], h <- [0 .. heads - 1], j <- [0 .. width - 1]]
    ([Int], [Int], [Int])
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"multihead attention equal divisible feature dimensions")

-- | Build a query-by-key allowed mask, with an explicit cached-prefix offset.
causalMask :: Int -> Int -> Int -> Either MathError [[Bool]]
causalMask :: Int -> Int -> Int -> Either MathError [[Bool]]
causalMask Int
queryLength Int
keyLength Int
offset
    | Int
queryLength Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
keyLength Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
        MathError -> Either MathError [[Bool]]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"causal mask lengths/offset")
    | Bool
otherwise =
        [[Bool]] -> Either MathError [[Bool]]
forall a b. b -> Either a b
Right
            [ [Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
j Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset | Int
j <- [Int
0 .. Int
keyLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            | Int
i <- [Int
0 .. Int
queryLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]

-- | Rotate adjacent feature pairs at sequence positions offset+i; feature dimension must be even.
rotaryEmbedding :: Double -> Int -> T.Tensor Double -> Either MathError (T.Tensor Double)
rotaryEmbedding :: Double -> Int -> Tensor Double -> Either MathError (Tensor Double)
rotaryEmbedding Double
base Int
offset Tensor Double
input = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
input
    [Double] -> Either MathError ()
finite [Double
base]
    case Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input of
        [Int
n, Int
d]
            | Double
base Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1 Bool -> Bool -> Bool
&& Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int -> Bool
forall a. Integral a => a -> Bool
even Int
d ->
                [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake
                    [Int
n, Int
d]
                    [ let pair :: Int
pair = Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                          angle :: Double
angle = Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
base Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
pair Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d)
                          a :: Double
a = Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int
i, Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
pair]
                          b :: Double
b = Tensor Double -> [Int] -> Double
forall a. Tensor a -> [Int] -> a
get Tensor Double
input [Int
i, Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
pair Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1]
                       in if Int -> Bool
forall a. Integral a => a -> Bool
even Int
j then Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos Double
angle Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin Double
angle else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin Double
angle Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos Double
angle
                    | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    , Int
j <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
        [Int]
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"rotary embedding base/offset/even feature dimension")

-- | Return [position,feature] sine/cosine positional encoding with base 10000.
sinusoidalEncoding :: Int -> Int -> Int -> Either MathError (T.Tensor Double)
sinusoidalEncoding :: Int -> Int -> Int -> Either MathError (Tensor Double)
sinusoidalEncoding Int
count Int
features Int
offset
    | Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
features Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sinusoidal encoding dimensions")
    | Bool
otherwise =
        [Int] -> [Double] -> Either MathError (Tensor Double)
checkedMake
            [Int
count, Int
features]
            [ let angle :: Double
angle =
                    Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset)
                        Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
10000 Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
features)
               in if Int -> Bool
forall a. Integral a => a -> Bool
even Int
j then Double -> Double
forall a. Floating a => a -> a
sin Double
angle else Double -> Double
forall a. Floating a => a -> a
cos Double
angle
            | Int
i <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            , Int
j <- [Int
0 .. Int
features Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]

-- | Return nonpositive slope-scaled absolute position offsets; combine with a causal allowed mask if required.
alibiBias :: Double -> Int -> Int -> Int -> Either MathError [[Double]]
alibiBias :: Double -> Int -> Int -> Int -> Either MathError [[Double]]
alibiBias Double
slope Int
queryLength Int
keyLength Int
offset = do
    [Double] -> Either MathError ()
finite [Double
slope]
    if Double
slope Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Int
queryLength Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
keyLength Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"ALiBi parameters")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let values :: [[Double]]
values =
            [ [ -Double
slope Double -> Double -> Double
forall a. Num a => a -> a -> a
* Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
forall a. Num a => a -> a
abs (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
j))
              | Int
j <- [Int
0 .. Int
keyLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
              ]
            | Int
i <- [Int
0 .. Int
queryLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]
    [Double] -> Either MathError ()
finite ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
values)
    [[Double]] -> Either MathError [[Double]]
forall a b. b -> Either a b
Right [[Double]]
values

{- | Opaque retained attention keys/values for incremental decoding; use checked cache
append/attention operations.
-}
data KVCache = KVCache Int Int Int (T.Tensor Double) (T.Tensor 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)

-- | Create an owned immutable key/value cache with fixed positive capacity and feature widths.
emptyCache :: Int -> Int -> Int -> Either MathError KVCache
emptyCache :: Int -> Int -> Int -> Either MathError KVCache
emptyCache Int
capacity Int
keyFeatures Int
valueFeatures
    | Int
capacity Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
keyFeatures Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
valueFeatures Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 =
        MathError -> Either MathError KVCache
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"cache capacity/features")
    | Bool
otherwise = do
        k <- [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int
0, Int
keyFeatures] []
        v <- T.tensor [0, valueFeatures] []
        Right (KVCache capacity keyFeatures valueFeatures k v)

-- | Append equal-length key/value sequences without truncating history; reject capacity or feature mismatches.
appendCache :: KVCache -> T.Tensor Double -> T.Tensor Double -> Either MathError KVCache
appendCache :: KVCache
-> Tensor Double -> Tensor Double -> Either MathError KVCache
appendCache (KVCache Int
capacity Int
kd Int
vd Tensor Double
keys Tensor Double
values) Tensor Double
k Tensor Double
v = do
    Tensor Double -> Either MathError ()
validTensor Tensor Double
k
    Tensor Double -> Either MathError ()
validTensor Tensor Double
v
    case (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
k, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
v, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
keys) of
        ([Int
n, Int
d], [Int
m, Int
f], [Int
old, Int
_]) | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
kd Bool -> Bool -> Bool
&& Int
f Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
vd Bool -> Bool -> Bool
&& Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
old Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
capacity -> do
            ks <- [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int
old Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n, Int
kd] (Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
keys [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
k)
            vs <- T.tensor [old + n, vd] (T.tensorValues values ++ T.tensorValues v)
            Right (KVCache capacity kd vd ks vs)
        ([Int], [Int], [Int])
_ -> MathError -> Either MathError KVCache
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"cache append shape or capacity")

-- | Return immutable accumulated key/value tensors, preserving insertion order.
cacheValues :: KVCache -> (T.Tensor Double, T.Tensor Double)
cacheValues :: KVCache -> (Tensor Double, Tensor Double)
cacheValues (KVCache Int
_ Int
_ Int
_ Tensor Double
k Tensor Double
v) = (Tensor Double
k, Tensor Double
v)

-- | Compose dense, checked elementwise activation, and dense layers with explicit parameters.
transformerFeedForward ::
    Activation ->
    T.Tensor Double ->
    T.Tensor Double ->
    [Double] ->
    T.Tensor Double ->
    [Double] ->
    Either MathError (T.Tensor Double)
transformerFeedForward :: Activation
-> Tensor Double
-> Tensor Double
-> [Double]
-> Tensor Double
-> [Double]
-> Either MathError (Tensor Double)
transformerFeedForward Activation
kind Tensor Double
input Tensor Double
w1 [Double]
b1 Tensor Double
w2 [Double]
b2 = do
    first <- Tensor Double
-> Tensor Double -> [Double] -> Either MathError (Tensor Double)
dense Tensor Double
input Tensor Double
w1 [Double]
b1
    values <- mapM (activation kind) (T.tensorValues first)
    middle <- T.tensor (T.tensorShape first) values
    dense middle w2 b2

{- | Symbolic trailing-axis layer normalization. Affine gamma/beta can be
composed with broadcast multiplication/addition in HMath.Autodiff.
-}
layerNormAD :: Int -> Double -> AD.ADTensor -> Either MathError AD.ADTensor
layerNormAD :: Int -> Double -> ADTensor -> Either MathError ADTensor
layerNormAD = Bool -> Int -> Double -> ADTensor -> Either MathError ADTensor
normalizeAD Bool
False

-- | Symbolic RMS normalization; its denominator includes all trailing axes.
rmsNormAD :: Int -> Double -> AD.ADTensor -> Either MathError AD.ADTensor
rmsNormAD :: Int -> Double -> ADTensor -> Either MathError ADTensor
rmsNormAD = Bool -> Int -> Double -> ADTensor -> Either MathError ADTensor
normalizeAD Bool
True

normalizeAD :: Bool -> Int -> Double -> AD.ADTensor -> Either MathError AD.ADTensor
normalizeAD :: Bool -> Int -> Double -> ADTensor -> Either MathError ADTensor
normalizeAD Bool
rms Int
axes Double
eps ADTensor
input = do
    let shape :: [Int]
shape = ADTensor -> [Int]
AD.adShape ADTensor
input; values :: [Expr]
values = ADTensor -> [Expr]
AD.adValues ADTensor
input; prefix :: [Int]
prefix = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
axes) [Int]
shape
    [Double] -> Either MathError ()
finite [Double
eps]
    if Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
axes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Bool -> Bool -> Bool
|| Double
eps Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| [Expr] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Expr]
values
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"symbolic normalization axes/epsilon/empty input")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let group :: [Int] -> Int
group [Int]
ix = [Int] -> [Int] -> Int
flat [Int]
prefix (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
prefix) [Int]
ix)
        allIx :: [[Int]]
allIx = [Int] -> [[Int]]
indices [Int]
shape
        normalized :: [Int] -> Expr
normalized [Int]
ix =
            let members :: [Expr]
members = [Expr
v | ([Int]
key, Expr
v) <- [[Int]] -> [Expr] -> [([Int], Expr)]
forall a b. [a] -> [b] -> [(a, b)]
zip [[Int]]
allIx [Expr]
values, [Int] -> Int
group [Int]
key Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
group [Int]
ix]
                size :: Double
size = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
members)
                mean :: Expr
mean = if Bool
rms then Expr
0 else [Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Expr]
members Expr -> Expr -> Expr
forall a. Fractional a => a -> a -> a
/ Double -> Expr
AD.constant Double
size
                variance :: Expr
variance = [Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Expr
v Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
mean) Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* (Expr
v Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
mean) | Expr
v <- [Expr]
members] Expr -> Expr -> Expr
forall a. Fractional a => a -> a -> a
/ Double -> Expr
AD.constant Double
size
             in ([Expr]
values [Expr] -> Int -> Expr
forall a. HasCallStack => [a] -> Int -> a
!! [Int] -> [Int] -> Int
flat [Int]
shape [Int]
ix Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
mean) Expr -> Expr -> Expr
forall a. Fractional a => a -> a -> a
/ Expr -> Expr
forall a. Floating a => a -> a
sqrt (Expr
variance Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
+ Double -> Expr
AD.constant Double
eps)
    [Int] -> [Expr] -> Either MathError ADTensor
AD.adTensor [Int]
shape (([Int] -> Expr) -> [[Int]] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Expr
normalized [[Int]]
allIx)

{- | First-order reverse attention graph, including Q/K score and V gradients.
Masks have the same True=allowed semantics as the checked numeric operation;
a fully masked row is a constant zero graph and backpropagates exact zeros.
-}
scaledDotProductAttentionAD ::
    AD.ADTensor -> AD.ADTensor -> AD.ADTensor -> Maybe [[Bool]] -> Either MathError AD.ADTensor
scaledDotProductAttentionAD :: ADTensor
-> ADTensor
-> ADTensor
-> Maybe [[Bool]]
-> Either MathError ADTensor
scaledDotProductAttentionAD ADTensor
query ADTensor
key ADTensor
value Maybe [[Bool]]
mask = case (ADTensor -> [Int]
AD.adShape ADTensor
query, ADTensor -> [Int]
AD.adShape ADTensor
key, ADTensor -> [Int]
AD.adShape ADTensor
value) of
    ([Int
n, Int
d], [Int
m, Int
k], [Int
v, Int
features]) | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k Bool -> Bool -> Bool
&& Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
        let allowed :: [[Bool]]
allowed = [[Bool]] -> ([[Bool]] -> [[Bool]]) -> Maybe [[Bool]] -> [[Bool]]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> [Bool] -> [[Bool]]
forall a. Int -> a -> [a]
replicate Int
n (Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Int
m Bool
True)) [[Bool]] -> [[Bool]]
forall a. a -> a
id Maybe [[Bool]]
mask
            q :: [Expr]
q = ADTensor -> [Expr]
AD.adValues ADTensor
query
            ks :: [Expr]
ks = ADTensor -> [Expr]
AD.adValues ADTensor
key
            vs :: [Expr]
vs = ADTensor -> [Expr]
AD.adValues ADTensor
value
        if [[Bool]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Bool]]
allowed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n Bool -> Bool -> Bool
|| ([Bool] -> Bool) -> [[Bool]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m) (Int -> Bool) -> ([Bool] -> Int) -> [Bool] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Bool]]
allowed
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"symbolic attention mask shape")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let row :: Int -> [Expr]
row Int
i =
                let active :: [Int]
active = [Int
j | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], ([[Bool]]
allowed [[Bool]] -> Int -> [Bool]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) [Bool] -> Int -> Bool
forall a. HasCallStack => [a] -> Int -> a
!! Int
j]
                    scores :: [Expr]
scores =
                        [ [Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Expr]
q [Expr] -> Int -> Expr
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t) Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* [Expr]
ks [Expr] -> Int -> Expr
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t) | Int
t <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Expr -> Expr -> Expr
forall a. Fractional a => a -> a -> a
/ Double -> Expr
AD.constant (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d))
                        | Int
j <- [Int]
active
                        ]
                    probabilities :: [Expr]
probabilities = [Expr] -> [Expr]
softmaxExpr [Expr]
scores
                 in [[Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Expr
p Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* [Expr]
vs [Expr] -> Int -> Expr
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
features Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t) | (Int
j, Expr
p) <- [Int] -> [Expr] -> [(Int, Expr)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
active [Expr]
probabilities] | Int
t <- [Int
0 .. Int
features Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        [Int] -> [Expr] -> Either MathError ADTensor
AD.adTensor [Int
n, Int
features] ((Int -> [Expr]) -> [Int] -> [Expr]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Int -> [Expr]
row [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
    ([Int], [Int], [Int])
_ -> MathError -> Either MathError ADTensor
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"symbolic attention shape")

checkedDimensions :: [Integer] -> Either MathError [Int]
checkedDimensions :: [Integer] -> Either MathError [Int]
checkedDimensions [Integer]
dimensions
    | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Integer
d -> Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)) [Integer]
dimensions =
        MathError -> Either MathError [Int]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"neural layer output dimension overflow")
    | Bool
otherwise = [Int] -> Either MathError [Int]
forall a b. b -> Either a b
Right ((Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
dimensions)