-- SPDX-License-Identifier: MPL-2.0

{- | Checked Double activations, losses, and real-coordinate reverse rules.
Probabilities must be in [0,1]; impossible events return a domain failure
instead of silently clipping probabilities. Logit losses avoid exp overflow.
Nondifferentiable hinges use zero at equality; MAE uses signum(0)=0.
-}
module HMath.DeepLearning.Core (
    Activation (..),
    activation,
    activationDerivative,
    activationExpr,
    softmax,
    logSoftmax,
    sparsemax,
    gatedLinearUnit,
    softmaxExpr,
    logSoftmaxExpr,
    sparsemaxExpr,
    ctcLossExpr,
    Reduction (..),
    Loss (..),
    loss,
    lossGradient,
    lossExpr,
    binaryCrossEntropyLogits,
    categoricalCrossEntropyLogits,
    regularization,
    contrastiveLoss,
    tripletLoss,
    focalLoss,
    diceLoss,
    iouLoss,
    ctcLoss,
    ctcLossGradient,
) where

import Data.List (sortBy)
import qualified HMath.Autodiff as AD
import HMath.Error

{- | Scalar activation selector. LeakyReLU/PReLU carry negative-side slope and ELU carries
positive alpha; other constructors select the named standard activation.
-}
data Activation
    = Sigmoid
    | Tanh
    | ReLU
    | LeakyReLU Double
    | PReLU Double
    | ELU Double
    | SELU
    | GELU
    | SiLU
    | Mish
    | Softplus
    | Softsign
    | HardSigmoid
    | HardSwish
    deriving (Activation -> Activation -> Bool
(Activation -> Activation -> Bool)
-> (Activation -> Activation -> Bool) -> Eq Activation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Activation -> Activation -> Bool
== :: Activation -> Activation -> Bool
$c/= :: Activation -> Activation -> Bool
/= :: Activation -> Activation -> Bool
Eq, Int -> Activation -> ShowS
[Activation] -> ShowS
Activation -> String
(Int -> Activation -> ShowS)
-> (Activation -> String)
-> ([Activation] -> ShowS)
-> Show Activation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Activation -> ShowS
showsPrec :: Int -> Activation -> ShowS
$cshow :: Activation -> String
show :: Activation -> String
$cshowList :: [Activation] -> ShowS
showList :: [Activation] -> 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
sigmoid :: Double -> Double
sigmoid :: Double -> Double
sigmoid Double
x = if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then 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 -> Double
forall a. Floating a => a -> a
exp (-Double
x)) else let z :: Double
z = Double -> Double
forall a. Floating a => a -> a
exp Double
x in Double
z Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
z)
softplus :: Double -> Double
softplus :: Double -> Double
softplus Double
x = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
x Double
0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (-Double -> Double
forall a. Num a => a -> a
abs Double
x))

-- GELU uses the explicit tanh approximation from Hendrycks/Gimpel (2016).

{- | Given Activation selector then finite scalar, evaluate the selected activation.
Parameterized LeakyReLU/PReLU slopes must be finite and ELU alpha positive. Invalid
input/parameters or non-finite output returns Left MathError.
-}
activation :: Activation -> Double -> Either MathError Double
activation :: Activation -> Double -> Either MathError Double
activation Activation
kind Double
x = do
    [Double] -> Either MathError ()
finite [Double
x]
    case Activation
kind of
        LeakyReLU Double
a -> [Double] -> Either MathError ()
finite [Double
a]
        PReLU Double
a -> [Double] -> Either MathError ()
finite [Double
a]
        ELU Double
a -> do
            [Double] -> Either MathError ()
finite [Double
a]
            if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive ELU alpha required") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        Activation
_ -> () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let y :: Double
y = case Activation
kind of
            Activation
Sigmoid -> Double -> Double
sigmoid Double
x
            Activation
Tanh -> Double -> Double
forall a. Floating a => a -> a
tanh Double
x
            Activation
ReLU -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 Double
x
            LeakyReLU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
x else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
            PReLU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
x else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
            ELU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
x else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Floating a => a -> a
exp Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1)
            Activation
SELU -> Double
1.0507009873554805 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
x else Double
1.6732632423543772 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Floating a => a -> a
exp Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1))
            Activation
GELU -> Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
tanh (Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.044715 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)))
            Activation
SiLU -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
sigmoid Double
x
            Activation
Mish -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
tanh (Double -> Double
softplus Double
x)
            Activation
Softplus -> Double -> Double
softplus Double
x
            Activation
Softsign -> Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Num a => a -> a
abs Double
x)
            Activation
HardSigmoid -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
6 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5))
            Activation
HardSwish -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
6 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5))
    [Double] -> Either MathError ()
finite [Double
y]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
y

-- | Evaluate the analytic first derivative with documented kink conventions and stable saturated sigmoid/tanh behavior.
activationDerivative :: Activation -> Double -> Either MathError Double
activationDerivative :: Activation -> Double -> Either MathError Double
activationDerivative Activation
kind Double
x = do
    _ <- Activation -> Double -> Either MathError Double
activation Activation
kind Double
x
    let s = Double -> Double
sigmoid Double
x
        t = Double -> Double
forall a. Floating a => a -> a
tanh Double
x
        y = case Activation
kind of
            Activation
Sigmoid -> Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s)
            Activation
Tanh -> Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t
            Activation
ReLU -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
1 else Double
0
            LeakyReLU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
1 else Double
a
            PReLU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
1 else Double
a
            ELU Double
a -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
1 else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp Double
x
            Activation
SELU -> Double
1.0507009873554805 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
1 else Double
1.6732632423543772 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp Double
x)
            Activation
GELU
                | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
20 -> Double
1
                | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< -Double
20 -> Double
0
                | Bool
otherwise ->
                    let c :: Double
c = Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi); u :: Double
u = Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.044715 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x); z :: Double
z = Double -> Double
forall a. Floating a => a -> a
tanh Double
u
                     in Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
z) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
0.044715 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)
            Activation
SiLU -> Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s)
            Activation
Mish -> let z :: Double
z = Double -> Double
forall a. Floating a => a -> a
tanh (Double -> Double
softplus Double
x) in Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s
            Activation
Softplus -> Double
s
            Activation
Softsign -> 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 -> Double
forall a. Num a => a -> a
abs Double
x) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)
            Activation
HardSigmoid -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> -Double
3 Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
3 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
6 else Double
0
            Activation
HardSwish -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= -Double
3 then Double
0 else if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
3 then Double
1 else Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5
    finite [y]
    Right y

-- | Embed the checked activation and analytic custom reverse rule in an automatic differentiation graph.
activationExpr :: Activation -> AD.Expr -> AD.Expr
activationExpr :: Activation -> Expr -> Expr
activationExpr Activation
a Expr
x = ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation [Double] -> Either MathError (Double, [Double])
rule [Expr
x]
  where
    rule :: [Double] -> Either MathError (Double, [Double])
rule [Double
v] = do y <- Activation -> Double -> Either MathError Double
activation Activation
a Double
v; dy <- activationDerivative a v; Right (y, [dy])
    rule [Double]
_ = MathError -> Either MathError (Double, [Double])
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"activation arity")

-- | Normalize a nonempty finite vector in log space after subtracting its maximum.
logSoftmax :: [Double] -> Either MathError [Double]
logSoftmax :: [Double] -> Either MathError [Double]
logSoftmax [] = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"empty softmax axis")
logSoftmax [Double]
xs = do
    [Double] -> Either MathError ()
finite [Double]
xs
    let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs
        normalizer :: Double
normalizer = Double -> Double
forall a. Floating a => a -> a
log ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double -> Double
forall a. Floating a => a -> a
exp (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) | Double
x <- [Double]
xs])
        ys :: [Double]
ys = (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
m) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
normalizer) [Double]
xs
    [Double] -> Either MathError ()
finite [Double]
ys
    [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Double]
ys

-- | Exponentiate stable log-softmax values; the output sums to one within floating error.
softmax :: [Double] -> Either MathError [Double]
softmax :: [Double] -> Either MathError [Double]
softmax [Double]
xs = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Floating a => a -> a
exp ([Double] -> [Double])
-> Either MathError [Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double] -> Either MathError [Double]
logSoftmax [Double]
xs

-- | Project a finite vector onto the probability simplex using sorted active-support thresholds.
sparsemax :: [Double] -> Either MathError [Double]
sparsemax :: [Double] -> Either MathError [Double]
sparsemax [] = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"empty sparsemax axis")
sparsemax [Double]
xs = do
    [Double] -> Either MathError ()
finite [Double]
xs
    let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs
        shifted :: [Double]
shifted = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
m) [Double]
xs
        sorted :: [Double]
sorted = (Double -> Double -> Ordering) -> [Double] -> [Double]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Double -> Double -> Ordering) -> Double -> Double -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip Double -> Double -> Ordering
forall a. Ord a => a -> a -> Ordering
compare) [Double]
shifted
        prefix :: [Double]
prefix = [Double] -> [Double]
forall {a}. [a] -> [a]
tailSafe ((Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 [Double]
sorted)
        active :: [(Int, Double)]
active =
            [ (Int
rank, Double
total)
            | (Int
rank, (Double
z, Double
total)) <- [Int] -> [(Double, Double)] -> [(Int, (Double, Double))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] ([Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
sorted [Double]
prefix)
            , Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rank Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
total
            ]
        (Int
k, Double
s) = (Int, Double) -> [(Int, Double)] -> (Int, Double)
forall {t}. t -> [t] -> t
lastSafe (Int
1, [Double]
sorted [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [(Int, Double)]
active
        threshold :: Double
threshold = (Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
        ys :: [Double]
ys = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
threshold)) [Double]
shifted
    [Double] -> Either MathError ()
finite [Double]
ys
    [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Double]
ys
  where
    tailSafe :: [a] -> [a]
tailSafe [] = []; tailSafe (a
_ : [a]
ys) = [a]
ys
    lastSafe :: t -> [t] -> t
lastSafe t
z [] = t
z; lastSafe t
_ [t
y] = t
y; lastSafe t
z (t
_ : [t]
ys) = t -> [t] -> t
lastSafe t
z [t]
ys

-- | Multiply values by sigmoid(gates), requiring equal vector lengths.
gatedLinearUnit :: [Double] -> [Double] -> Either MathError [Double]
gatedLinearUnit :: [Double] -> [Double] -> Either MathError [Double]
gatedLinearUnit [Double]
values [Double]
gate = do
    [Double] -> Either MathError ()
finite ([Double]
values [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
gate)
    if [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
gate
        then MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values) ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
gate))
        else [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
x Double
g -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
sigmoid Double
g) [Double]
values [Double]
gate)

{- | Loss aggregation: Sum adds elementwise losses and Mean divides by their count;
vector-level loss exceptions are documented by loss.
-}
data Reduction = Sum | Mean deriving (Reduction -> Reduction -> Bool
(Reduction -> Reduction -> Bool)
-> (Reduction -> Reduction -> Bool) -> Eq Reduction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Reduction -> Reduction -> Bool
== :: Reduction -> Reduction -> Bool
$c/= :: Reduction -> Reduction -> Bool
/= :: Reduction -> Reduction -> Bool
Eq, Int -> Reduction -> ShowS
[Reduction] -> ShowS
Reduction -> String
(Int -> Reduction -> ShowS)
-> (Reduction -> String)
-> ([Reduction] -> ShowS)
-> Show Reduction
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Reduction -> ShowS
showsPrec :: Int -> Reduction -> ShowS
$cshow :: Reduction -> String
show :: Reduction -> String
$cshowList :: [Reduction] -> ShowS
showList :: [Reduction] -> ShowS
Show)

{- | Checked loss selector,including squared/absolute error,cross
entropy,KL,hinge,cosine,and parameterized Huber/smooth losses. Target/prediction domain
rules belong to loss.
-}
data Loss
    = MeanSquaredError
    | MeanAbsoluteError
    | Huber Double
    | SmoothL1 Double
    | BinaryCrossEntropy
    | CategoricalCrossEntropy
    | NegativeLogLikelihood
    | KLDivergence
    | CosineLoss
    | Hinge Double
    deriving (Loss -> Loss -> Bool
(Loss -> Loss -> Bool) -> (Loss -> Loss -> Bool) -> Eq Loss
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Loss -> Loss -> Bool
== :: Loss -> Loss -> Bool
$c/= :: Loss -> Loss -> Bool
/= :: Loss -> Loss -> Bool
Eq, Int -> Loss -> ShowS
[Loss] -> ShowS
Loss -> String
(Int -> Loss -> ShowS)
-> (Loss -> String) -> ([Loss] -> ShowS) -> Show Loss
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Loss -> ShowS
showsPrec :: Int -> Loss -> ShowS
$cshow :: Loss -> String
show :: Loss -> String
$cshowList :: [Loss] -> ShowS
showList :: [Loss] -> ShowS
Show)

validatePair :: [Double] -> [Double] -> Either MathError ()
validatePair :: [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
prediction = do
    [Double] -> Either MathError ()
finite ([Double]
target [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
prediction)
    if [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
target Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
target Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
prediction
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"nonempty equal loss vectors required")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
probability :: [Double] -> Either MathError ()
probability :: [Double] -> Either MathError ()
probability [Double]
xs =
    if (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1) [Double]
xs then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"probability outside [0,1]") else [Double] -> Either MathError ()
finite [Double]
xs

-- | Evaluate the selected checked loss. Mean divides elementwise losses by count; CosineLoss is one vector loss.
loss :: Loss -> Reduction -> [Double] -> [Double] -> Either MathError Double
loss :: Loss
-> Reduction -> [Double] -> [Double] -> Either MathError Double
loss Loss
kind Reduction
reduction [Double]
target [Double]
prediction = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
prediction
    value <- case Loss
kind of
        Loss
CosineLoss -> do
            (a, _, _) <- [Double] -> Either MathError ([Double], Double, Double)
unitVector [Double]
target
            (b, _, _) <- unitVector prediction
            Right (1 - sum (zipWith (*) a b))
        Loss
_ -> do
            vals <- ((Double, Double) -> Either MathError Double)
-> [(Double, Double)] -> 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 (Double, Double) -> Either MathError Double
term ([Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
target [Double]
prediction)
            Right (sum vals / (if reduction == Mean then fromIntegral (length vals) else 1))
    finite [value]
    Right value
  where
    term :: (Double, Double) -> Either MathError Double
term (Double
t, Double
p) = case Loss
kind of
        Loss
MeanSquaredError -> Double -> Either MathError Double
forall a b. b -> Either a b
Right ((Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int))
        Loss
MeanAbsoluteError -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t))
        Huber Double
delta
            | Double
delta Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 ->
                Double -> Either MathError Double
forall a b. b -> Either a b
Right (if Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
delta then (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 else Double
delta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
delta Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2))
            | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive Huber delta required")
        SmoothL1 Double
beta
            | Double
beta Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 ->
                Double -> Either MathError Double
forall a b. b -> Either a b
Right (if Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
beta then (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
beta) else Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
            | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive smooth L1 beta required")
        Loss
BinaryCrossEntropy -> do [Double] -> Either MathError ()
probability [Double
t, Double
p]; Double -> Double -> Either MathError Double
forall {b}. (Ord b, Floating b) => b -> b -> Either MathError b
logTerm Double
t Double
p Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Double
a -> Double -> Double -> Either MathError Double
forall {b}. (Ord b, Floating b) => b -> b -> Either MathError b
logTerm (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p) Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Double
b -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)
        Loss
CategoricalCrossEntropy -> do [Double] -> Either MathError ()
probability [Double
t, Double
p]; Double -> Double -> Either MathError Double
forall {b}. (Ord b, Floating b) => b -> b -> Either MathError b
logTerm Double
t Double
p
        Loss
NegativeLogLikelihood -> do
            [Double] -> Either MathError ()
probability [Double
t]
            if Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"NLL expects nonpositive log probabilities") else Double -> Either MathError Double
forall a b. b -> Either a b
Right (-Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p)
        Loss
KLDivergence -> do
            [Double] -> Either MathError ()
probability [Double
t, Double
p]
            if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
0
                else if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"KL zero denominator") else Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p))
        Hinge Double
margin
            | Double
margin Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double
margin Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p))
            | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"hinge margin or +/-1 label")
        Loss
CosineLoss -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"internal cosine loss")
    logTerm :: b -> b -> Either MathError b
logTerm b
0 b
_ = b -> Either MathError b
forall a b. b -> Either a b
Right b
0
    logTerm b
t b
p
        | b
p b -> b -> Bool
forall a. Ord a => a -> a -> Bool
<= b
0 = MathError -> Either MathError b
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive probability required for nonzero target")
        | Bool
otherwise = b -> Either MathError b
forall a b. b -> Either a b
Right (-b
t b -> b -> b
forall a. Num a => a -> a -> a
* b -> b
forall a. Floating a => a -> a
log b
p)

-- | Analytic derivatives with respect to predictions under the exact loss/reduction convention.
lossGradient :: Loss -> Reduction -> [Double] -> [Double] -> Either MathError [Double]
lossGradient :: Loss
-> Reduction -> [Double] -> [Double] -> Either MathError [Double]
lossGradient Loss
kind Reduction
reduction [Double]
target [Double]
prediction = do
    _ <- Loss
-> Reduction -> [Double] -> [Double] -> Either MathError Double
loss Loss
kind Reduction
reduction [Double]
target [Double]
prediction
    let n = if Reduction
reduction Reduction -> Reduction -> Bool
forall a. Eq a => a -> a -> Bool
== Reduction
Mean Bool -> Bool -> Bool
&& Loss
kind Loss -> Loss -> Bool
forall a. Eq a => a -> a -> Bool
/= Loss
CosineLoss then 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]
target) else Double
1
        derivatives = case Loss
kind of
            Loss
CosineLoss -> case ([Double] -> Either MathError ([Double], Double, Double)
unitVector [Double]
target, [Double] -> Either MathError ([Double], Double, Double)
unitVector [Double]
prediction) of
                (Right ([Double]
a, Double
_, Double
_), Right ([Double]
b, Double
m, Double
s)) -> let cosine :: Double
cosine = [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]
a [Double]
b) in (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
t Double
p -> ((Double
cosine Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) [Double]
a [Double]
b
                (Either MathError ([Double], Double, Double),
 Either MathError ([Double], Double, Double))
_ -> Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
target) (Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
            Loss
_ -> (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
derivative [Double]
target [Double]
prediction
    finite derivatives
    Right (map (/ n) derivatives)
  where
    derivative :: Double -> Double -> Double
derivative Double
t Double
p = case Loss
kind of
        Loss
MeanSquaredError -> Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t)
        Loss
MeanAbsoluteError -> Double -> Double
forall a. Num a => a -> a
signum (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t)
        Huber Double
d -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (-Double
d) (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
d (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t))
        SmoothL1 Double
b -> if Double -> Double
forall a. Num a => a -> a
abs (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
b then (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
b else Double -> Double
forall a. Num a => a -> a
signum (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t)
        Loss
BinaryCrossEntropy -> (if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else -Double
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then Double
0 else (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p))
        Loss
CategoricalCrossEntropy -> if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else -Double
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p
        Loss
NegativeLogLikelihood -> -Double
t
        Loss
KLDivergence -> if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else -Double
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p
        Hinge Double
margin -> if Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
margin then -Double
t else Double
0
        Loss
CosineLoss -> Double
0

-- | Use a constant target with symbolic prediction inputs and the matching analytic reverse rule.
lossExpr :: Loss -> Reduction -> [Double] -> [AD.Expr] -> AD.Expr
lossExpr :: Loss -> Reduction -> [Double] -> [Expr] -> Expr
lossExpr Loss
kind Reduction
reduction [Double]
target =
    ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation
        ( \[Double]
xs -> do y <- Loss
-> Reduction -> [Double] -> [Double] -> Either MathError Double
loss Loss
kind Reduction
reduction [Double]
target [Double]
xs; gs <- lossGradient kind reduction target xs; Right (y, gs)
        )

-- | Stable binary cross entropy from logits and [0,1] targets, including extreme finite logits.
binaryCrossEntropyLogits :: Reduction -> [Double] -> [Double] -> Either MathError Double
binaryCrossEntropyLogits :: Reduction -> [Double] -> [Double] -> Either MathError Double
binaryCrossEntropyLogits Reduction
reduction [Double]
target [Double]
logits = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
logits
    [Double] -> Either MathError ()
probability [Double]
target
    let values :: [Double]
values = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
t Double
x -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
x Double
0 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (-Double -> Double
forall a. Num a => a -> a
abs Double
x))) [Double]
target [Double]
logits
        y :: Double
y = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
values Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (if Reduction
reduction Reduction -> Reduction -> Bool
forall a. Eq a => a -> a -> Bool
== Reduction
Mean then 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]
values) else Double
1)
    [Double] -> Either MathError ()
finite [Double
y]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
y

-- | Cross entropy from one class-logit vector; nonnegative targets must sum to one.
categoricalCrossEntropyLogits :: [Double] -> [Double] -> Either MathError Double
categoricalCrossEntropyLogits :: [Double] -> [Double] -> Either MathError Double
categoricalCrossEntropyLogits [Double]
target [Double]
logits = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
logits
    [Double] -> Either MathError ()
probability [Double]
target
    if Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
target Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1e-10
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"categorical targets must sum to one")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    logs <- [Double] -> Either MathError [Double]
logSoftmax [Double]
logits
    Right (-sum (zipWith (*) target logs))

-- | Return l1*sum(abs(w)) + l2/2*sum(w*w); both coefficients must be finite and nonnegative.
regularization :: Double -> Double -> [Double] -> Either MathError Double
regularization :: Double -> Double -> [Double] -> Either MathError Double
regularization Double
l1 Double
l2 [Double]
weights = do
    [Double] -> Either MathError ()
finite (Double
l1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Double
l2 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
weights)
    if Double
l1 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
l2 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"negative regularization") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let value :: Double
value =
            (if Double
l1 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
l1 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]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
weights))
                Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (if Double
l2 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
l2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 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]
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]
weights))
    [Double] -> Either MathError ()
finite [Double
value]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
value

-- | Pairwise Euclidean contrastive loss; similar=True penalizes squared distance and dissimilar pairs use the positive margin.
contrastiveLoss :: Double -> Bool -> [Double] -> [Double] -> Either MathError Double
contrastiveLoss :: Double -> Bool -> [Double] -> [Double] -> Either MathError Double
contrastiveLoss Double
margin Bool
similar [Double]
a [Double]
b = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
a [Double]
b
    [Double] -> Either MathError ()
finite [Double
margin]
    if Double
margin Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive contrastive margin") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let d2 :: Double
d2 = [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
x Double
y -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
a [Double]
b)
    let value :: Double
value = if Bool
similar then Double
d2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 else Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double
margin Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
sqrt Double
d2) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
    [Double] -> Either MathError ()
finite [Double
value]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
value

-- | Return max(0,positiveDistance-negativeDistance+margin) for equal nonempty vector dimensions.
tripletLoss :: Double -> [Double] -> [Double] -> [Double] -> Either MathError Double
tripletLoss :: Double
-> [Double] -> [Double] -> [Double] -> Either MathError Double
tripletLoss Double
margin [Double]
anchor [Double]
positive [Double]
negative = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
anchor [Double]
positive
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
anchor [Double]
negative
    [Double] -> Either MathError ()
finite [Double
margin]
    if Double
margin Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"negative triplet margin") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let distance :: [a] -> [a] -> a
distance [a]
a [a]
b = a -> a
forall a. Floating a => a -> a
sqrt ([a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
x a
y -> (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
y) a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [a]
a [a]
b))
    let value :: Double
value = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 ([Double] -> [Double] -> Double
forall {a}. Floating a => [a] -> [a] -> a
distance [Double]
anchor [Double]
positive Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> [Double] -> Double
forall {a}. Floating a => [a] -> [a] -> a
distance [Double]
anchor [Double]
negative Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
margin)
    [Double] -> Either MathError ()
finite [[Double] -> [Double] -> Double
forall {a}. Floating a => [a] -> [a] -> a
distance [Double]
anchor [Double]
positive, [Double] -> [Double] -> Double
forall {a}. Floating a => [a] -> [a] -> a
distance [Double]
anchor [Double]
negative, Double
value]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
value

-- | Mean binary focal loss with alpha in [0,1], gamma>=0, and exact zero/one targets.
focalLoss :: Double -> Double -> [Double] -> [Double] -> Either MathError Double
focalLoss :: Double -> Double -> [Double] -> [Double] -> Either MathError Double
focalLoss Double
alpha Double
gamma [Double]
target [Double]
prediction = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
prediction
    [Double] -> Either MathError ()
probability (Double
alpha Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
target [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
prediction)
    [Double] -> Either MathError ()
finite [Double
gamma]
    if Double
gamma Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
t -> Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0 Bool -> Bool -> Bool
&& Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
1) [Double]
target
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"focal gamma or binary labels")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    vals <-
        ((Double, Double) -> Either MathError Double)
-> [(Double, Double)] -> 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
            ( \(Double
t, Double
p) ->
                let pt :: Double
pt = if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then Double
p else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p; at :: Double
at = if Double
t Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then Double
alpha else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
alpha
                 in if Double
pt Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                        then MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"focal impossible event")
                        else Double -> Either MathError Double
forall a b. b -> Either a b
Right (-Double
at Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
pt) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
pt)
            )
            ([Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
target [Double]
prediction)
    Right (sum vals / fromIntegral (length vals))

-- | Soft segmentation overlap losses on [0,1] values; positive smoothing defines the empty-mask case.
diceLoss, iouLoss :: Double -> [Double] -> [Double] -> Either MathError Double
diceLoss :: Double -> [Double] -> [Double] -> Either MathError Double
diceLoss Double
smooth [Double]
target [Double]
prediction = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
prediction
    [Double] -> Either MathError ()
probability ([Double]
target [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
prediction)
    [Double] -> Either MathError ()
finite [Double
smooth]
    if Double
smooth Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive Dice smoothing") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    Double -> Either MathError Double
forall a b. b -> Either a b
Right
        (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
2 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]
target [Double]
prediction) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
smooth) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
target 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]
prediction Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
smooth))
iouLoss :: Double -> [Double] -> [Double] -> Either MathError Double
iouLoss Double
smooth [Double]
target [Double]
prediction = do
    [Double] -> [Double] -> Either MathError ()
validatePair [Double]
target [Double]
prediction
    [Double] -> Either MathError ()
probability ([Double]
target [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
prediction)
    [Double] -> Either MathError ()
finite [Double
smooth]
    if Double
smooth Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive IoU smoothing") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let intersection :: Double
intersection = [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]
target [Double]
prediction)
    Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
intersection Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
smooth) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
target 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]
prediction Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
intersection Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
smooth))

{- | Single sequence, time-major log-probabilities [time,class]. Blank index is
explicit; labels exclude blank. The forward/backward dynamic program sums
alignments in log space, including adjacent repeated-label restrictions.
Gradients are with respect to independent log probabilities, before a
log-softmax chain rule. Impossible alignments return DomainError.
-}
ctcLoss :: Int -> [Int] -> [[Double]] -> Either MathError Double
ctcLoss :: Int -> [Int] -> [[Double]] -> Either MathError Double
ctcLoss Int
blank [Int]
labels [[Double]]
logs = (Double, [[Double]]) -> Double
forall a b. (a, b) -> a
fst ((Double, [[Double]]) -> Double)
-> Either MathError (Double, [[Double]]) -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Int] -> [[Double]] -> Either MathError (Double, [[Double]])
ctcLossGradient Int
blank [Int]
labels [[Double]]
logs

{- | Given blank class index, target labels, then rectangular time-by-class log
probabilities, return CTC negative log likelihood and gradient rows. Validate class
indices, finite/log-probability constraints and alignment feasibility. Empty time
permits only empty labels; invalid/impossible inputs return Left MathError.
-}
ctcLossGradient :: Int -> [Int] -> [[Double]] -> Either MathError (Double, [[Double]])
ctcLossGradient :: Int -> [Int] -> [[Double]] -> Either MathError (Double, [[Double]])
ctcLossGradient Int
blank [Int]
labels [[Double]]
logs = do
    case [[Double]]
logs of
        [] -> if [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
labels then (Double, [[Double]]) -> Either MathError (Double, [[Double]])
forall a b. b -> Either a b
Right (Double
0, []) else MathError -> Either MathError (Double, [[Double]])
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"CTC nonempty labels with empty time")
        [Double]
first : [[Double]]
_ ->
            if [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
first 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
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
first) (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]]
logs
                then MathError -> Either MathError (Double, [[Double]])
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"CTC rectangular class dimension")
                else Int -> Either MathError (Double, [[Double]])
calculate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
first)
  where
    calculate :: Int -> Either MathError (Double, [[Double]])
calculate Int
classes = do
        [Double] -> Either MathError ()
finite ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
logs)
        if Int
blank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            Bool -> Bool -> Bool
|| Int
blank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
classes
            Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\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
classes Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blank) [Int]
labels
            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) ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
logs)
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"CTC blank, labels, or log probabilities")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let states :: [Int]
states = (Int -> [Int]) -> [Int] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Int
l -> [Int
blank, Int
l]) [Int]
labels [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
blank]
            count :: Int
count = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
states
            times :: Int
times = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
logs
            neg :: Double
neg = -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0
            initial :: [Double]
initial = [if Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then ([[Double]]
logs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s) else Double
neg | Int
s <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            forwardStep :: [a] -> [a] -> [a]
forwardStep [a]
previous [a]
row =
                [ [a]
row [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s)
                    a -> a -> a
forall a. Num a => a -> a -> a
+ [a] -> a
forall {t}. RealFloat t => [t] -> t
logAdd
                        [ [a]
previous [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                        | Int
j <- [Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) .. Int
s]
                        , Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Bool -> Bool -> Bool
|| ([Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
blank Bool -> Bool -> Bool
&& [Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
j)
                        ]
                | Int
s <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
            alphas :: [[Double]]
alphas = ([Double] -> [Double] -> [Double])
-> [Double] -> [[Double]] -> [[Double]]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl [Double] -> [Double] -> [Double]
forall {a}. RealFloat a => [a] -> [a] -> [a]
forwardStep [Double]
initial (Int -> [[Double]] -> [[Double]]
forall a. Int -> [a] -> [a]
drop Int
1 [[Double]]
logs)
            terminal :: [Double]
terminal = [if Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 then Double
0 else Double
neg | Int
s <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            backwardStep :: (Int, [Double]) -> p -> (Int, [Double])
backwardStep (Int
time, [Double]
next) p
_ =
                let row :: [Double]
row = [[Double]]
logs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! (Int
time Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    prev :: [Double]
prev =
                        [ [Double] -> Double
forall {t}. RealFloat t => [t] -> t
logAdd
                            [ [Double]
row [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
next [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                            | Int
j <- [Int
s .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)]
                            , Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Bool -> Bool -> Bool
|| ([Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
blank Bool -> Bool -> Bool
&& [Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s)
                            ]
                        | Int
s <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        ]
                 in (Int
time Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, [Double]
prev)
            backwards :: [[Double]]
backwards = ((Int, [Double]) -> [Double]) -> [(Int, [Double])] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (Int, [Double]) -> [Double]
forall a b. (a, b) -> b
snd (((Int, [Double]) -> Int -> (Int, [Double]))
-> (Int, [Double]) -> [Int] -> [(Int, [Double])]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (Int, [Double]) -> Int -> (Int, [Double])
forall {p}. (Int, [Double]) -> p -> (Int, [Double])
backwardStep (Int
times Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2, [Double]
terminal) [Int
1 .. Int
times Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
            betas :: [[Double]]
betas = [[Double]] -> [[Double]]
forall {a}. [a] -> [a]
reverse [[Double]]
backwards
            total :: Double
total =
                [Double] -> Double
forall {t}. RealFloat t => [t] -> t
logAdd [[Double]
lastRow [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
s | Int
s <- [Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], let lastRow :: [Double]
lastRow = [[Double]]
alphas [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! (Int
times Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
        if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
total Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
total
            then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"CTC alignment has zero probability")
            else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
        let gradients :: [[Double]]
gradients =
                [ [ -[Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double -> Double
forall a. Floating a => a -> a
exp (([[Double]]
alphas [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
t) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ([[Double]]
betas [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
t) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
total) | Int
s <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
states [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
c]
                  | Int
c <- [Int
0 .. Int
classes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                  ]
                | Int
t <- [Int
0 .. Int
times 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]]
gradients)
        (Double, [[Double]]) -> Either MathError (Double, [[Double]])
forall a b. b -> Either a b
Right (-Double
total, [[Double]]
gradients)
    logAdd :: [t] -> t
logAdd [t]
xs = let m :: t
m = [t] -> t
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [t]
xs in if t -> Bool
forall a. RealFloat a => a -> Bool
isInfinite t
m Bool -> Bool -> Bool
&& t
m t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
0 then t
m else t
m t -> t -> t
forall a. Num a => a -> a -> a
+ t -> t
forall a. Floating a => a -> a
log ([t] -> t
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [t -> t
forall a. Floating a => a -> a
exp (t
x t -> t -> t
forall a. Num a => a -> a -> a
- t
m) | t
x <- [t]
xs])

-- | First-order softmax reverse rule, including cross-coordinate coupling.
softmaxExpr :: [AD.Expr] -> [AD.Expr]
softmaxExpr :: [Expr] -> [Expr]
softmaxExpr [Expr]
xs = [([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation (Int -> [Double] -> Either MathError (Double, [Double])
rule Int
i) [Expr]
xs | Int
i <- [Int
0 .. [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    rule :: Int -> [Double] -> Either MathError (Double, [Double])
rule Int
i [Double]
values = do
        ps <- [Double] -> Either MathError [Double]
softmax [Double]
values
        let p = [Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
        Right (p, [p * ((if i == j then 1 else 0) - q) | (j, q) <- zip [0 :: Int ..] ps])

-- | First-order stable log-softmax rule; gradient rows sum to zero.
logSoftmaxExpr :: [AD.Expr] -> [AD.Expr]
logSoftmaxExpr :: [Expr] -> [Expr]
logSoftmaxExpr [Expr]
xs = [([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation (Int -> [Double] -> Either MathError (Double, [Double])
rule Int
i) [Expr]
xs | Int
i <- [Int
0 .. [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    rule :: Int -> [Double] -> Either MathError (Double, [Double])
rule Int
i [Double]
values = do
        logs <- [Double] -> Either MathError [Double]
logSoftmax [Double]
values
        Right (logs !! i, [(if i == j then 1 else 0) - exp p | (j, p) <- zip [0 :: Int ..] logs])

{- | Sparsemax's piecewise-linear Jacobian. At a support boundary its inactive
coordinate convention chooses zero, giving a defined subgradient.
-}
sparsemaxExpr :: [AD.Expr] -> [AD.Expr]
sparsemaxExpr :: [Expr] -> [Expr]
sparsemaxExpr [Expr]
xs = [([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation (Int -> [Double] -> Either MathError (Double, [Double])
forall {a}.
Fractional a =>
Int -> [Double] -> Either MathError (Double, [a])
rule Int
i) [Expr]
xs | Int
i <- [Int
0 .. [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    rule :: Int -> [Double] -> Either MathError (Double, [a])
rule Int
i [Double]
values = do
        ps <- [Double] -> Either MathError [Double]
sparsemax [Double]
values
        let count = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
filter (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) [Double]
ps)
            active = [Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
        Right
            ( ps !! i
            , [ if active && p > 0 then (if i == j then 1 else 0) - 1 / fromIntegral count else 0
              | (j, p) <- zip [0 :: Int ..] ps
              ]
            )

{- | CTC graph operation on rectangular time-major log probabilities. Its custom
backward uses forward/backward occupancies, rather than numerical differences.
-}
ctcLossExpr :: Int -> [Int] -> [[AD.Expr]] -> Either MathError AD.Expr
ctcLossExpr :: Int -> [Int] -> [[Expr]] -> Either MathError Expr
ctcLossExpr Int
blank [Int]
labels [[Expr]]
logs = case [[Expr]]
logs of
    [] -> if [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
labels then Expr -> Either MathError Expr
forall a b. b -> Either a b
Right Expr
0 else MathError -> Either MathError Expr
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"empty CTC time")
    [Expr]
first : [[Expr]]
_
        | [Expr] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Expr]
first Bool -> Bool -> Bool
|| ([Expr] -> Bool) -> [[Expr]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
first) (Int -> Bool) -> ([Expr] -> Int) -> [Expr] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Expr]]
logs ->
            MathError -> Either MathError Expr
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"CTC symbolic rectangular shape")
        | Bool
otherwise ->
            let classes :: Int
classes = [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
first; times :: Int
times = [[Expr]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Expr]]
logs
             in Expr -> Either MathError Expr
forall a b. b -> Either a b
Right
                    ( ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
AD.customOperation
                        ( \[Double]
values -> do
                            (value, grads) <-
                                Int -> [Int] -> [[Double]] -> Either MathError (Double, [[Double]])
ctcLossGradient Int
blank [Int]
labels [Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
classes (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
classes) [Double]
values) | Int
t <- [Int
0 .. Int
times Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                            Right (value, concat grads)
                        )
                        ([[Expr]] -> [Expr]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Expr]]
logs)
                    )

unitVector :: [Double] -> Either MathError ([Double], Double, Double)
unitVector :: [Double] -> Either MathError ([Double], Double, Double)
unitVector [Double]
values = do
    let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
values)
    if Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"cosine loss zero vector") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let scaled :: [Double]
scaled = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) [Double]
values
        s :: Double
s = 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]
scaled))
    ([Double], Double, Double)
-> Either MathError ([Double], Double, Double)
forall a b. b -> Either a b
Right ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) [Double]
scaled, Double
m, Double
s)