{-# LANGUAGE FlexibleInstances #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Pure CPU automatic differentiation. 'Dual' nests for higher orders. 'Expr'
compiles to an immutable, reusable reverse tape; dynamic graphs are ordinary
Haskell functions building expressions at run time. Tapes own their saved
primals, so no ambient mutable graph or lifetime is involved. Accumulation is
explicit in 'GradientState'; reverse evaluation itself never accumulates.
Custom rules return a value and local partial derivatives (first order only).
Checkpoints store the subgraph recipe and recompute its tape on backprop.
Complex gradients use real coordinates: dL/dz = (dL/dx+i*dL/dy)/2 (Wirtinger).
-}
module HMath.Autodiff (
    Dual (..),
    primal,
    tangent,
    forward,
    forwardJVP,
    Expr,
    variable,
    constant,
    customOperation,
    customUnary,
    detach,
    stopGradient,
    checkpoint,
    CompiledGraph,
    compileGraph,
    evaluateGraph,
    gradient,
    valueAndGradient,
    differentiate,
    hessian,
    jacobian,
    jvp,
    vjp,
    hvp,
    GradientState (..),
    zeroGradients,
    accumulateGradients,
    resetGradients,
    gradientHook,
    clipGradient,
    gradientCheck,
    complexGradient,
    differentiableStep,
    ADTensor,
    adTensor,
    adValues,
    adShape,
    adBroadcast,
    adAdd,
    adMultiply,
    adMatmul,
    adContract,
    adSum,
) where

import Control.Monad (foldM)
import Data.Complex (Complex ((:+)))
import qualified Data.Foldable as Foldable (foldl')
import qualified Data.IntMap.Strict as IM
import HMath.Error
import qualified HMath.Tensor as T

-- | A tangent may itself be a dual number. Branches use the primal ordering.
data Dual a = Dual a a deriving (Dual a -> Dual a -> Bool
(Dual a -> Dual a -> Bool)
-> (Dual a -> Dual a -> Bool) -> Eq (Dual a)
forall a. Eq a => Dual a -> Dual a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Dual a -> Dual a -> Bool
== :: Dual a -> Dual a -> Bool
$c/= :: forall a. Eq a => Dual a -> Dual a -> Bool
/= :: Dual a -> Dual a -> Bool
Eq, Int -> Dual a -> ShowS
[Dual a] -> ShowS
Dual a -> String
(Int -> Dual a -> ShowS)
-> (Dual a -> String) -> ([Dual a] -> ShowS) -> Show (Dual a)
forall a. Show a => Int -> Dual a -> ShowS
forall a. Show a => [Dual a] -> ShowS
forall a. Show a => Dual a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Dual a -> ShowS
showsPrec :: Int -> Dual a -> ShowS
$cshow :: forall a. Show a => Dual a -> String
show :: Dual a -> String
$cshowList :: forall a. Show a => [Dual a] -> ShowS
showList :: [Dual a] -> ShowS
Show)

primal, tangent :: Dual a -> a

{- | Extract the primal value from Dual value tangent without evaluating a derivative.
Example: primal (Dual 3 1) == 3.
-}
primal :: forall a. Dual a -> a
primal (Dual a
a a
_) = a
a

{- | Extract the directional derivative seed/result from a Dual value. Example: tangent
(Dual 3 1) == 1.
-}
tangent :: forall a. Dual a -> a
tangent (Dual a
_ a
a) = a
a

instance (Ord a) => Ord (Dual a) where compare :: Dual a -> Dual a -> Ordering
compare (Dual a
x a
_) (Dual a
y a
_) = a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y
instance (Num a) => Num (Dual a) where
    Dual a
x a
dx + :: Dual a -> Dual a -> Dual a
+ Dual a
y a
dy = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
y) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
+ a
dy)
    Dual a
x a
dx - :: Dual a -> Dual a -> Dual a
- Dual a
y a
dy = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
y) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
- a
dy)
    Dual a
x a
dx * :: Dual a -> Dual a -> Dual a
* Dual a
y a
dy = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
y) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
dy)
    negate :: Dual a -> Dual a
negate (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Num a => a -> a
negate a
x) (a -> a
forall a. Num a => a -> a
negate a
dx)
    abs :: Dual a -> Dual a
abs (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Num a => a -> a
abs a
x) (a -> a
forall a. Num a => a -> a
signum a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
dx)
    signum :: Dual a -> Dual a
signum (Dual a
x a
_) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Num a => a -> a
signum a
x) a
0
    fromInteger :: Integer -> Dual a
fromInteger Integer
n = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
n) a
0
instance (Fractional a) => Fractional (Dual a) where
    recip :: Dual a -> Dual a
recip (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Fractional a => a -> a
recip a
x) (-a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))
    fromRational :: Rational -> Dual a
fromRational Rational
r = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (Rational -> a
forall a. Fractional a => Rational -> a
fromRational Rational
r) a
0
instance (Floating a, Eq a) => Floating (Dual a) where
    Dual a
x a
dx ** :: Dual a -> Dual a -> Dual a
** Dual a
y a
dy
        | a
dy a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
y) (if a
y a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
0 else a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Floating a => a -> a -> a
** (a
y a -> a -> a
forall a. Num a => a -> a -> a
- a
1))
        | Bool
otherwise = let value :: a
value = a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
y in a -> a -> Dual a
forall a. a -> a -> Dual a
Dual a
value (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Floating a => a -> a -> a
** (a
y a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Num a => a -> a -> a
+ a
dy a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
log a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
value)
    pi :: Dual a
pi = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual a
forall a. Floating a => a
pi a
0
    exp :: Dual a -> Dual a
exp (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
exp a
x) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
exp a
x)
    log :: Dual a -> Dual a
log (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
log a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
x)
    sqrt :: Dual a -> Dual a
sqrt (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
sqrt a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sqrt a
x))
    sin :: Dual a -> Dual a
sin (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
sin a
x) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
cos a
x)
    cos :: Dual a -> Dual a
cos (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
cos a
x) (-a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sin a
x)
    tan :: Dual a -> Dual a
tan (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
tan a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a -> a
forall a. Floating a => a -> a
cos a
x a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
cos a
x))
    asin :: Dual a -> Dual a
asin (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
asin a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
sqrt (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))
    acos :: Dual a -> Dual a
acos (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
acos a
x) (-a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
sqrt (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))
    atan :: Dual a -> Dual a
atan (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
atan a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))
    sinh :: Dual a -> Dual a
sinh (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
sinh a
x) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
cosh a
x)
    cosh :: Dual a -> Dual a
cosh (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
cosh a
x) (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sinh a
x)
    tanh :: Dual a -> Dual a
tanh (Dual a
x a
dx) = let t :: a
t = a -> a
forall a. Floating a => a -> a
tanh a
x in a -> a -> Dual a
forall a. a -> a -> Dual a
Dual a
t (a
dx a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
t a -> a -> a
forall a. Num a => a -> a -> a
* a
t))
    asinh :: Dual a -> Dual a
asinh (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
asinh a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
sqrt (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))
    acosh :: Dual a -> Dual a
acosh (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
acosh a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
sqrt (a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
1))
    atanh :: Dual a -> Dual a
atanh (Dual a
x a
dx) = a -> a -> Dual a
forall a. a -> a -> Dual a
Dual (a -> a
forall a. Floating a => a -> a
atanh a
x) (a
dx a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
x))

-- | Evaluate a scalar dual function with unit tangent; returns primal and first derivative.
forward :: (Num a) => (Dual a -> Dual a) -> a -> (a, a)
forward :: forall a. Num a => (Dual a -> Dual a) -> a -> (a, a)
forward Dual a -> Dual a
f a
x = let Dual a
y a
dy = Dual a -> Dual a
f (a -> a -> Dual a
forall a. a -> a -> Dual a
Dual a
x a
1) in (a
y, a
dy)

-- | Evaluate a vector dual function along a checked same-length direction; every output must be finite.
forwardJVP ::
    ([Dual Double] -> [Dual Double]) -> [Double] -> [Double] -> Either MathError ([Double], [Double])
forwardJVP :: ([Dual Double] -> [Dual Double])
-> [Double] -> [Double] -> Either MathError ([Double], [Double])
forwardJVP [Dual Double] -> [Dual Double]
f [Double]
xs [Double]
ds = do
    [Double] -> [Double] -> Either MathError ()
forall a b. [a] -> [b] -> Either MathError ()
same [Double]
xs [Double]
ds
    [Double] -> Either MathError ()
finite [Double]
xs
    [Double] -> Either MathError ()
finite [Double]
ds
    let ys :: [Dual Double]
ys = [Dual Double] -> [Dual Double]
f ((Double -> Double -> Dual Double)
-> [Double] -> [Double] -> [Dual Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Dual Double
forall a. a -> a -> Dual a
Dual [Double]
xs [Double]
ds)
    [Double] -> Either MathError ()
finite ((Dual Double -> Double) -> [Dual Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Dual Double -> Double
forall a. Dual a -> a
primal [Dual Double]
ys [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ (Dual Double -> Double) -> [Dual Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Dual Double -> Double
forall a. Dual a -> a
tangent [Dual Double]
ys)
    ([Double], [Double]) -> Either MathError ([Double], [Double])
forall a b. b -> Either a b
Right ((Dual Double -> Double) -> [Dual Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Dual Double -> Double
forall a. Dual a -> a
primal [Dual Double]
ys, (Dual Double -> Double) -> [Dual Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Dual Double -> Double
forall a. Dual a -> a
tangent [Dual Double]
ys)

data Unary
    = Power Double
    | Exp
    | Log
    | Sqrt
    | Sin
    | Cos
    | Tan
    | Asin
    | Acos
    | Atan
    | Sinh
    | Cosh
    | Tanh
    | Asinh
    | Acosh
    | Atanh
    | Abs
    | Sign
    deriving (Unary -> Unary -> Bool
(Unary -> Unary -> Bool) -> (Unary -> Unary -> Bool) -> Eq Unary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Unary -> Unary -> Bool
== :: Unary -> Unary -> Bool
$c/= :: Unary -> Unary -> Bool
/= :: Unary -> Unary -> Bool
Eq, Int -> Unary -> ShowS
[Unary] -> ShowS
Unary -> String
(Int -> Unary -> ShowS)
-> (Unary -> String) -> ([Unary] -> ShowS) -> Show Unary
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Unary -> ShowS
showsPrec :: Int -> Unary -> ShowS
$cshow :: Unary -> String
show :: Unary -> String
$cshowList :: [Unary] -> ShowS
showList :: [Unary] -> ShowS
Show)

{- | Opaque reverse-mode differentiable scalar expression,composed from explicit inputs and
operations. Compile/evaluate through the checked autodiff graph API.
-}
data Expr
    = C Double
    | V Int
    | Add Expr Expr
    | Mul Expr Expr
    | Neg Expr
    | Inv Expr
    | U Unary Expr
    | Stop Expr
    | Custom ([Double] -> Either MathError (Double, [Double])) [Expr]
    | Checkpoint ([Expr] -> Expr) [Expr]

instance Show Expr where
    show :: Expr -> String
show (C Double
x) = Double -> String
forall a. Show a => a -> String
show Double
x
    show (V Int
i) = String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i
    show (Add Expr
a Expr
b) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"+" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
b String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Mul Expr
a Expr
b) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"*" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
b String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Neg Expr
a) = String
"-" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a
    show (Inv Expr
a) = String
"1/" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a
    show (U Unary
op Expr
a) = Unary -> String
forall a. Show a => a -> String
show Unary
op String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Stop Expr
a) = String
"detach(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. Show a => a -> String
show Expr
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show (Custom [Double] -> Either MathError (Double, [Double])
_ [Expr]
_) = String
"custom-operation"
    show (Checkpoint [Expr] -> Expr
_ [Expr]
_) = String
"checkpoint"
instance Num Expr where
    + :: Expr -> Expr -> Expr
(+) = Expr -> Expr -> Expr
Add
    * :: Expr -> Expr -> Expr
(*) = Expr -> Expr -> Expr
Mul
    negate :: Expr -> Expr
negate = Expr -> Expr
Neg
    abs :: Expr -> Expr
abs = Unary -> Expr -> Expr
U Unary
Abs
    signum :: Expr -> Expr
signum = Unary -> Expr -> Expr
U Unary
Sign
    fromInteger :: Integer -> Expr
fromInteger = Double -> Expr
C (Double -> Expr) -> (Integer -> Double) -> Integer -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Double
forall a. Num a => Integer -> a
fromInteger
instance Fractional Expr where
    recip :: Expr -> Expr
recip = Expr -> Expr
Inv
    fromRational :: Rational -> Expr
fromRational = Double -> Expr
C (Double -> Expr) -> (Rational -> Double) -> Rational -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> Double
forall a. Fractional a => Rational -> a
fromRational
instance Floating Expr where
    Expr
base ** :: Expr -> Expr -> Expr
** C Double
powerValue = Unary -> Expr -> Expr
U (Double -> Unary
Power Double
powerValue) Expr
base
    Expr
base ** Expr
powerValue = Expr -> Expr
forall a. Floating a => a -> a
exp (Expr
powerValue Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr -> Expr
forall a. Floating a => a -> a
log Expr
base)
    pi :: Expr
pi = Double -> Expr
C Double
forall a. Floating a => a
pi
    exp :: Expr -> Expr
exp = Unary -> Expr -> Expr
U Unary
Exp
    log :: Expr -> Expr
log = Unary -> Expr -> Expr
U Unary
Log
    sqrt :: Expr -> Expr
sqrt = Unary -> Expr -> Expr
U Unary
Sqrt
    sin :: Expr -> Expr
sin = Unary -> Expr -> Expr
U Unary
Sin
    cos :: Expr -> Expr
cos = Unary -> Expr -> Expr
U Unary
Cos
    tan :: Expr -> Expr
tan = Unary -> Expr -> Expr
U Unary
Tan
    asin :: Expr -> Expr
asin = Unary -> Expr -> Expr
U Unary
Asin
    acos :: Expr -> Expr
acos = Unary -> Expr -> Expr
U Unary
Acos
    atan :: Expr -> Expr
atan = Unary -> Expr -> Expr
U Unary
Atan
    sinh :: Expr -> Expr
sinh = Unary -> Expr -> Expr
U Unary
Sinh
    cosh :: Expr -> Expr
cosh = Unary -> Expr -> Expr
U Unary
Cosh
    tanh :: Expr -> Expr
tanh = Unary -> Expr -> Expr
U Unary
Tanh
    asinh :: Expr -> Expr
asinh = Unary -> Expr -> Expr
U Unary
Asinh
    acosh :: Expr -> Expr
acosh = Unary -> Expr -> Expr
U Unary
Acosh
    atanh :: Expr -> Expr
atanh = Unary -> Expr -> Expr
U Unary
Atanh

-- | Declare a zero-based graph input. compileGraph validates the declared input count.
variable :: Int -> Expr
variable :: Int -> Expr
variable = Int -> Expr
V

-- | Embed a literal; nonfinite literals are rejected when a graph is compiled.
constant :: Double -> Expr
constant :: Double -> Expr
constant = Double -> Expr
C

-- | Supply one local partial derivative per argument. Rules must be deterministic and total; outputs and arity are checked.
customOperation :: ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
customOperation :: ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
customOperation = ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
Custom

-- | First-order scalar custom rule; the second callback is the analytic local derivative.
customUnary :: (Double -> Double) -> (Double -> Double) -> Expr -> Expr
customUnary :: (Double -> Double) -> (Double -> Double) -> Expr -> Expr
customUnary Double -> Double
f Double -> Double
df Expr
x = ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
Custom [Double] -> Either MathError (Double, [Double])
rule [Expr
x]
  where
    rule :: [Double] -> Either MathError (Double, [Double])
rule [Double
a] = (Double, [Double]) -> Either MathError (Double, [Double])
forall a b. b -> Either a b
Right (Double -> Double
f Double
a, [Double -> Double
df Double
a]); rule [Double]
_ = MathError -> Either MathError (Double, [Double])
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"custom unary arity")

-- | Preserve the primal value while removing every reverse edge through this expression.
detach, stopGradient :: Expr -> Expr
detach :: Expr -> Expr
detach = Expr -> Expr
Stop
stopGradient :: Expr -> Expr
stopGradient = Expr -> Expr
Stop

-- | Store a graph-building recipe instead of the internal tape; replay it during reverse propagation.
checkpoint :: ([Expr] -> Expr) -> [Expr] -> Expr
checkpoint :: ([Expr] -> Expr) -> [Expr] -> Expr
checkpoint = ([Expr] -> Expr) -> [Expr] -> Expr
Checkpoint

-- Nodes record only direct predecessor indices; a compiled graph can be reused
-- at new input values. Each invocation owns its fresh tape and adjoint map.
data Node
    = Literal Double
    | Input Int
    | SumNode Int Int
    | ProductNode Int Int
    | NegateNode Int
    | InverseNode Int
    | UnaryNode Unary Int
    | StopNode Int
    | CustomNode ([Double] -> Either MathError (Double, [Double])) [Int]
    | RecomputeNode ([Expr] -> Expr) [Int]

{- | Opaque reusable automatic-differentiation computation graph with explicit input/output
structure.
-}
data CompiledGraph = CompiledGraph Int Int [(Int, Node)]

data Pullback = Partials [(Int, Double)] | Deferred ([Expr] -> Expr) [Int] [Double]

-- | Compile an immutable reusable graph with n inputs and a hard limit of one million scalar nodes.
compileGraph :: Int -> Expr -> Either MathError CompiledGraph
compileGraph :: Int -> Expr -> Either MathError CompiledGraph
compileGraph Int
n Expr
e
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError CompiledGraph
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative graph input count")
    | Bool
otherwise = do
        (root, _, nodes) <- Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
build Int
0 [] Expr
e
        Right (CompiledGraph n root (reverse nodes))
  where
    add :: Int -> [(Int, b)] -> b -> Either MathError (Int, Int, [(Int, b)])
add Int
next [(Int, b)]
nodes b
node
        | Int
next Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1000000 = MathError -> Either MathError (Int, Int, [(Int, b)])
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"automatic differentiation graph node budget" Int
next)
        | Bool
otherwise = (Int, Int, [(Int, b)]) -> Either MathError (Int, Int, [(Int, b)])
forall a b. b -> Either a b
Right (Int
next, Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, (Int
next, b
node) (Int, b) -> [(Int, b)] -> [(Int, b)]
forall a. a -> [a] -> [a]
: [(Int, b)]
nodes)
    build :: Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
build Int
next [(Int, Node)]
nodes Expr
ex = case Expr
ex of
        C Double
x -> if Double -> Bool
isFinite Double
x then Int
-> [(Int, Node)]
-> Node
-> Either MathError (Int, Int, [(Int, Node)])
forall {b}.
Int -> [(Int, b)] -> b -> Either MathError (Int, Int, [(Int, b)])
add Int
next [(Int, Node)]
nodes (Double -> Node
Literal Double
x) else MathError -> Either MathError (Int, Int, [(Int, Node)])
forall a b. a -> Either a b
Left MathError
NonFiniteResult
        V Int
i -> if 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
n then Int
-> [(Int, Node)]
-> Node
-> Either MathError (Int, Int, [(Int, Node)])
forall {b}.
Int -> [(Int, b)] -> b -> Either MathError (Int, Int, [(Int, b)])
add Int
next [(Int, Node)]
nodes (Int -> Node
Input Int
i) else MathError -> Either MathError (Int, Int, [(Int, Node)])
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"graph input index")
        Add Expr
a Expr
b -> (Int -> Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
binary Int -> Int -> Node
SumNode Int
next [(Int, Node)]
nodes Expr
a Expr
b
        Mul Expr
a Expr
b -> (Int -> Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
binary Int -> Int -> Node
ProductNode Int
next [(Int, Node)]
nodes Expr
a Expr
b
        Neg Expr
a -> (Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
unary Int -> Node
NegateNode Int
next [(Int, Node)]
nodes Expr
a
        Inv Expr
a -> (Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
unary Int -> Node
InverseNode Int
next [(Int, Node)]
nodes Expr
a
        U (Power Double
powerValue) Expr
_ | Bool -> Bool
not (Double -> Bool
isFinite Double
powerValue) -> MathError -> Either MathError (Int, Int, [(Int, Node)])
forall a b. a -> Either a b
Left MathError
NonFiniteResult
        U Unary
op Expr
a -> (Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
unary (Unary -> Int -> Node
UnaryNode Unary
op) Int
next [(Int, Node)]
nodes Expr
a
        Stop Expr
a -> (Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
unary Int -> Node
StopNode Int
next [(Int, Node)]
nodes Expr
a
        Custom [Double] -> Either MathError (Double, [Double])
f [Expr]
es -> do (ids, k, ns) <- Int
-> [(Int, Node)]
-> [Expr]
-> Either MathError ([Int], Int, [(Int, Node)])
many Int
next [(Int, Node)]
nodes [Expr]
es; add k ns (CustomNode f ids)
        Checkpoint [Expr] -> Expr
f [Expr]
es -> do (ids, k, ns) <- Int
-> [(Int, Node)]
-> [Expr]
-> Either MathError ([Int], Int, [(Int, Node)])
many Int
next [(Int, Node)]
nodes [Expr]
es; add k ns (RecomputeNode f ids)
    unary :: (Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
unary Int -> Node
f Int
k [(Int, Node)]
ns Expr
a = do (i, j, ms) <- Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
build Int
k [(Int, Node)]
ns Expr
a; add j ms (f i)
    binary :: (Int -> Int -> Node)
-> Int
-> [(Int, Node)]
-> Expr
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
binary Int -> Int -> Node
f Int
k [(Int, Node)]
ns Expr
a Expr
b = do (i, j, ms) <- Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
build Int
k [(Int, Node)]
ns Expr
a; (l, q, rs) <- build j ms b; add q rs (f i l)
    many :: Int
-> [(Int, Node)]
-> [Expr]
-> Either MathError ([Int], Int, [(Int, Node)])
many Int
k [(Int, Node)]
ns [] = ([Int], Int, [(Int, Node)])
-> Either MathError ([Int], Int, [(Int, Node)])
forall a b. b -> Either a b
Right ([], Int
k, [(Int, Node)]
ns)
    many Int
k [(Int, Node)]
ns (Expr
a : [Expr]
as) = do (i, j, ms) <- Int
-> [(Int, Node)]
-> Expr
-> Either MathError (Int, Int, [(Int, Node)])
build Int
k [(Int, Node)]
ns Expr
a; (ids, q, rs) <- many j ms as; Right (i : ids, q, rs)

isFinite :: Double -> Bool
isFinite :: Double -> Bool
isFinite 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)
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 -> Bool
isFinite [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
same :: [a] -> [b] -> Either MathError ()
same :: forall a b. [a] -> [b] -> Either MathError ()
same [a]
a [b]
b = if [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
b then () -> Either MathError ()
forall a b. b -> Either a b
Right () else MathError -> Either MathError ()
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
a) ([b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
b))
unaryValue :: Unary -> Double -> Double
unaryValue :: Unary -> Double -> Double
unaryValue Unary
op = case Unary
op of
    Power Double
powerValue -> (Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
powerValue)
    Unary
Exp -> Double -> Double
forall a. Floating a => a -> a
exp
    Unary
Log -> Double -> Double
forall a. Floating a => a -> a
log
    Unary
Sqrt -> Double -> Double
forall a. Floating a => a -> a
sqrt
    Unary
Sin -> Double -> Double
forall a. Floating a => a -> a
sin
    Unary
Cos -> Double -> Double
forall a. Floating a => a -> a
cos
    Unary
Tan -> Double -> Double
forall a. Floating a => a -> a
tan
    Unary
Asin -> Double -> Double
forall a. Floating a => a -> a
asin
    Unary
Acos -> Double -> Double
forall a. Floating a => a -> a
acos
    Unary
Atan -> Double -> Double
forall a. Floating a => a -> a
atan
    Unary
Sinh -> Double -> Double
forall a. Floating a => a -> a
sinh
    Unary
Cosh -> Double -> Double
forall a. Floating a => a -> a
cosh
    Unary
Tanh -> Double -> Double
forall a. Floating a => a -> a
tanh
    Unary
Asinh -> Double -> Double
forall a. Floating a => a -> a
asinh
    Unary
Acosh -> Double -> Double
forall a. Floating a => a -> a
acosh
    Unary
Atanh -> Double -> Double
forall a. Floating a => a -> a
atanh
    Unary
Abs -> Double -> Double
forall a. Num a => a -> a
abs
    Unary
Sign -> Double -> Double
forall a. Num a => a -> a
signum

-- Expressions for local derivatives support exact symbolic higher orders.
localDerivative :: Unary -> Expr -> Expr
localDerivative :: Unary -> Expr -> Expr
localDerivative Unary
op Expr
x = case Unary
op of
    Power Double
powerValue -> if Double
powerValue Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Expr
0 else Double -> Expr
C Double
powerValue Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x Expr -> Expr -> Expr
forall a. Floating a => a -> a -> a
** Double -> Expr
C (Double
powerValue Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1)
    Unary
Exp -> Expr -> Expr
forall a. Floating a => a -> a
exp Expr
x
    Unary
Log -> Expr -> Expr
forall a. Fractional a => a -> a
recip Expr
x
    Unary
Sqrt -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr
2 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr -> Expr
forall a. Floating a => a -> a
sqrt Expr
x)
    Unary
Sin -> Expr -> Expr
forall a. Floating a => a -> a
cos Expr
x
    Unary
Cos -> Expr -> Expr
forall a. Num a => a -> a
negate (Expr -> Expr
forall a. Floating a => a -> a
sin Expr
x)
    Unary
Tan -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr -> Expr
forall a. Floating a => a -> a
cos Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr -> Expr
forall a. Floating a => a -> a
cos Expr
x)
    Unary
Asin -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr -> Expr
forall a. Floating a => a -> a
sqrt (Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x))
    Unary
Acos -> Expr -> Expr
forall a. Num a => a -> a
negate (Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr -> Expr
forall a. Floating a => a -> a
sqrt (Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x)))
    Unary
Atan -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
+ Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x)
    Unary
Sinh -> Expr -> Expr
forall a. Floating a => a -> a
cosh Expr
x
    Unary
Cosh -> Expr -> Expr
forall a. Floating a => a -> a
sinh Expr
x
    Unary
Tanh -> Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr -> Expr
forall a. Floating a => a -> a
tanh Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr -> Expr
forall a. Floating a => a -> a
tanh Expr
x
    Unary
Asinh -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr -> Expr
forall a. Floating a => a -> a
sqrt (Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
+ Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x))
    Unary
Acosh -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr -> Expr
forall a. Floating a => a -> a
sqrt (Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
1))
    Unary
Atanh -> Expr -> Expr
forall a. Fractional a => a -> a
recip (Expr
1 Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
x)
    Unary
Abs -> Expr -> Expr
forall a. Num a => a -> a
signum Expr
x
    Unary
Sign -> Expr
0

runTape ::
    CompiledGraph -> [Double] -> Either MathError (Double, [(Int, Pullback)], IM.IntMap Double)
runTape :: CompiledGraph
-> [Double]
-> Either MathError (Double, [(Int, Pullback)], IntMap Double)
runTape (CompiledGraph Int
n Int
root [(Int, Node)]
nodes) [Double]
xs = do
    if [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
n ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)) else [Double] -> Either MathError ()
finite [Double]
xs
    (vs, edges) <- (Either MathError (IntMap Double, [(Int, Pullback)])
 -> (Int, Node)
 -> Either MathError (IntMap Double, [(Int, Pullback)]))
-> Either MathError (IntMap Double, [(Int, Pullback)])
-> [(Int, Node)]
-> Either MathError (IntMap Double, [(Int, Pullback)])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Either MathError (IntMap Double, [(Int, Pullback)])
-> (Int, Node)
-> Either MathError (IntMap Double, [(Int, Pullback)])
step ((IntMap Double, [(Int, Pullback)])
-> Either MathError (IntMap Double, [(Int, Pullback)])
forall a b. b -> Either a b
Right (IntMap Double
forall a. IntMap a
IM.empty, [])) [(Int, Node)]
nodes
    Right (vs IM.! root, edges, vs)
  where
    step :: Either MathError (IntMap Double, [(Int, Pullback)])
-> (Int, Node)
-> Either MathError (IntMap Double, [(Int, Pullback)])
step Either MathError (IntMap Double, [(Int, Pullback)])
acc (Int
i, Node
node) = do
        (vs, es) <- Either MathError (IntMap Double, [(Int, Pullback)])
acc
        let get Int
j = IntMap Double
vs IntMap Double -> Int -> Double
forall a. IntMap a -> Int -> a
IM.! Int
j
        (v, ds) <- case node of
            Literal Double
x -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (Double
x, [])
            Input Int
k -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right ([Double]
xs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k, [])
            SumNode Int
a Int
b -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (Int -> Double
get Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
get Int
b, [(Int
a, Double
1), (Int
b, Double
1)])
            ProductNode Int
a Int
b -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (Int -> Double
get Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
get Int
b, [(Int
a, Int -> Double
get Int
b), (Int
b, Int -> Double
get Int
a)])
            NegateNode Int
a -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (-Int -> Double
get Int
a, [(Int
a, -Double
1)])
            InverseNode Int
a -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
get Int
a, [(Int
a, -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
get Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
get Int
a))])
            UnaryNode Unary
op Int
a -> do
                d <- Expr -> Either MathError Double
evaluate (Unary -> Expr -> Expr
localDerivative Unary
op (Double -> Expr
C (Int -> Double
get Int
a)))
                Right (unaryValue op (get a), [(a, d)])
            StopNode Int
a -> (Double, [(Int, Double)])
-> Either MathError (Double, [(Int, Double)])
forall a b. b -> Either a b
Right (Int -> Double
get Int
a, [])
            CustomNode [Double] -> Either MathError (Double, [Double])
f [Int]
ids -> do
                (v, partials) <- [Double] -> Either MathError (Double, [Double])
f ((Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
get [Int]
ids)
                same ids partials
                Right (v, zip ids partials)
            RecomputeNode [Expr] -> Expr
f [Int]
ids -> do
                graph <- Int -> Expr -> Either MathError CompiledGraph
compileGraph ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids) ([Expr] -> Expr
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))
                v <- evaluateGraph graph (map get ids)
                Right (v, [])
        finite (v : map snd ds)
        let edge = case Node
node of RecomputeNode [Expr] -> Expr
recipe [Int]
ids -> ([Expr] -> Expr) -> [Int] -> [Double] -> Pullback
Deferred [Expr] -> Expr
recipe [Int]
ids ((Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
get [Int]
ids); Node
_ -> [(Int, Double)] -> Pullback
Partials [(Int, Double)]
ds
        Right (IM.insert i v vs, (i, edge) : es)
    evaluate :: Expr -> Either MathError Double
evaluate Expr
e = do g <- Int -> Expr -> Either MathError CompiledGraph
compileGraph Int
0 Expr
e; (y, _, _) <- runValue g []; Right y

-- local derivative evaluation must not recursively differentiate its own
-- unary nodes. This separate evaluator computes values alone.
runValue :: CompiledGraph -> [Double] -> Either MathError (Double, (), ())
runValue :: CompiledGraph -> [Double] -> Either MathError (Double, (), ())
runValue (CompiledGraph Int
n Int
root [(Int, Node)]
ns) [Double]
xs = do
    if Int
n 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]
xs then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
n ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)) else [Double] -> Either MathError ()
finite [Double]
xs
    vs <- (Either MathError (IntMap Double)
 -> (Int, Node) -> Either MathError (IntMap Double))
-> Either MathError (IntMap Double)
-> [(Int, Node)]
-> Either MathError (IntMap 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' Either MathError (IntMap Double)
-> (Int, Node) -> Either MathError (IntMap Double)
step (IntMap Double -> Either MathError (IntMap Double)
forall a b. b -> Either a b
Right IntMap Double
forall a. IntMap a
IM.empty) [(Int, Node)]
ns
    Right (vs IM.! root, (), ())
  where
    step :: Either MathError (IntMap Double)
-> (Int, Node) -> Either MathError (IntMap Double)
step Either MathError (IntMap Double)
acc (Int
i, Node
node) = do
        vs <- Either MathError (IntMap Double)
acc
        let get Int
k = IntMap Double
vs IntMap Double -> Int -> Double
forall a. IntMap a -> Int -> a
IM.! Int
k
        y <- case node of
            Literal Double
x -> Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
x
            Input Int
k -> Double -> Either MathError Double
forall a b. b -> Either a b
Right ([Double]
xs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k)
            SumNode Int
a Int
b -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Int -> Double
get Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
get Int
b)
            ProductNode Int
a Int
b -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Int -> Double
get Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
get Int
b)
            NegateNode Int
a -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (-Int -> Double
get Int
a)
            InverseNode Int
a -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
get Int
a)
            UnaryNode Unary
op Int
a -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Unary -> Double -> Double
unaryValue Unary
op (Int -> Double
get Int
a))
            StopNode Int
a -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Int -> Double
get Int
a)
            CustomNode [Double] -> Either MathError (Double, [Double])
f [Int]
ids -> (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
<$> [Double] -> Either MathError (Double, [Double])
f ((Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
get [Int]
ids)
            RecomputeNode [Expr] -> Expr
f [Int]
ids -> do
                graph <- Int -> Expr -> Either MathError CompiledGraph
compileGraph ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids) ([Expr] -> Expr
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))
                evaluateGraph graph (map get ids)
        finite [y]
        Right (IM.insert i y vs)

-- | Re-evaluate a compiled graph at exactly n finite inputs, without requesting gradients.
evaluateGraph :: CompiledGraph -> [Double] -> Either MathError Double
evaluateGraph :: CompiledGraph -> [Double] -> Either MathError Double
evaluateGraph CompiledGraph
g [Double]
xs = do (y, _, _) <- CompiledGraph -> [Double] -> Either MathError (Double, (), ())
runValue CompiledGraph
g [Double]
xs; Right y

-- | Build a dynamic graph and return its checked value and reverse gradient; repeated input leaves accumulate.
valueAndGradient :: ([Expr] -> Expr) -> [Double] -> Either MathError (Double, [Double])
valueAndGradient :: ([Expr] -> Expr) -> [Double] -> Either MathError (Double, [Double])
valueAndGradient [Expr] -> Expr
f [Double]
xs = do
    let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
    g@(CompiledGraph _ root nodes) <- Int -> Expr -> Either MathError CompiledGraph
compileGraph Int
n ([Expr] -> Expr
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))
    (y, edges, _) <- runTape g xs
    adj <-
        foldM
            ( \IntMap Double
as (Int
i, Pullback
pullback) -> do
                derivatives <- case Pullback
pullback of
                    Partials [(Int, Double)]
ds -> [(Int, Double)] -> Either MathError [(Int, Double)]
forall a b. b -> Either a b
Right [(Int, Double)]
ds
                    Deferred [Expr] -> Expr
recipe [Int]
ids [Double]
primals -> do
                        (_, partials) <- ([Expr] -> Expr) -> [Double] -> Either MathError (Double, [Double])
valueAndGradient [Expr] -> Expr
recipe [Double]
primals
                        Right (zip ids partials)
                let adjoint = Double -> Int -> IntMap Double -> Double
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault Double
0 Int
i IntMap Double
as
                    next = (IntMap Double -> (Int, Double) -> IntMap Double)
-> IntMap Double -> [(Int, Double)] -> IntMap 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' (\IntMap Double
m (Int
j, Double
d) -> (Double -> Double -> Double)
-> Int -> Double -> IntMap Double -> IntMap Double
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IM.insertWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Int
j (Double
adjoint Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d) IntMap Double
m) IntMap Double
as [(Int, Double)]
derivatives
                finite (IM.elems next)
                Right next
            )
            (IM.singleton root 1)
            edges
    let grads =
            (IntMap Double -> (Int, Node) -> IntMap Double)
-> IntMap Double -> [(Int, Node)] -> IntMap 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'
                (\IntMap Double
m (Int
i, Node
node) -> case Node
node of Input Int
k -> (Double -> Double -> Double)
-> Int -> Double -> IntMap Double -> IntMap Double
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IM.insertWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Int
k (Double -> Int -> IntMap Double -> Double
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault Double
0 Int
i IntMap Double
adj) IntMap Double
m; Node
_ -> IntMap Double
m)
                IntMap Double
forall a. IntMap a
IM.empty
                [(Int, Node)]
nodes
        result = [Double -> Int -> IntMap Double -> Double
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault Double
0 Int
k IntMap Double
grads | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    finite result
    Right (y, result)

-- | Return a fresh reverse gradient. Previous calls and explicit GradientState values are unaffected.
gradient :: ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient :: ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient [Expr] -> Expr
f [Double]
xs = (Double, [Double]) -> [Double]
forall a b. (a, b) -> b
snd ((Double, [Double]) -> [Double])
-> Either MathError (Double, [Double]) -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Expr] -> Expr) -> [Double] -> Either MathError (Double, [Double])
valueAndGradient [Expr] -> Expr
f [Double]
xs

-- | Construct a symbolic derivative with respect to one graph input. Opaque custom first-order rules reject higher orders.
differentiate :: Int -> Expr -> Either MathError Expr
differentiate :: Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
ex = case Expr
ex of
    C Double
_ -> Expr -> Either MathError Expr
forall a b. b -> Either a b
Right Expr
0
    V Int
j -> Expr -> Either MathError Expr
forall a b. b -> Either a b
Right (if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then Expr
1 else Expr
0)
    Add Expr
a Expr
b -> Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
(+) (Expr -> Expr -> Expr)
-> Either MathError Expr -> Either MathError (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
a Either MathError (Expr -> Expr)
-> Either MathError Expr -> Either MathError Expr
forall a b.
Either MathError (a -> b)
-> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
b
    Mul Expr
a Expr
b -> do da <- Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
a; db <- differentiate i b; Right (da * b + a * db)
    Neg Expr
a -> Expr -> Expr
forall a. Num a => a -> a
negate (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
a
    Inv Expr
a -> (\Expr
da -> -Expr
da Expr -> Expr -> Expr
forall a. Fractional a => a -> a -> a
/ (Expr
a Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
a)) (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
a
    U Unary
op Expr
a -> (Unary -> Expr -> Expr
localDerivative Unary
op Expr
a Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
*) (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Expr -> Either MathError Expr
differentiate Int
i Expr
a
    Stop Expr
_ -> Expr -> Either MathError Expr
forall a b. b -> Either a b
Right Expr
0
    Custom [Double] -> Either MathError (Double, [Double])
_ [Expr]
_ ->
        MathError -> Either MathError Expr
forall a b. a -> Either a b
Left
            (String -> MathError
InvalidArgument String
"custom rules support first order; provide a symbolic operation for higher orders")
    Checkpoint [Expr] -> Expr
f [Expr]
as -> Int -> Expr -> Either MathError Expr
differentiate Int
i ([Expr] -> Expr
f [Expr]
as)

-- | Return output-by-input Jacobian rows, including shape-preserving empty output handling.
jacobian :: ([Expr] -> [Expr]) -> [Double] -> Either MathError [[Double]]
jacobian :: ([Expr] -> [Expr]) -> [Double] -> Either MathError [[Double]]
jacobian [Expr] -> [Expr]
f [Double]
xs = do
    [Double] -> Either MathError ()
finite [Double]
xs
    (Expr -> Either MathError [Double])
-> [Expr] -> 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 (\Expr
e -> ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient (Expr -> [Expr] -> Expr
forall a b. a -> b -> a
const Expr
e) [Double]
xs) ([Expr] -> [Expr]
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))

-- | Return input-by-input symbolic second derivatives; all evaluated derivatives must be finite.
hessian :: ([Expr] -> Expr) -> [Double] -> Either MathError [[Double]]
hessian :: ([Expr] -> Expr) -> [Double] -> Either MathError [[Double]]
hessian [Expr] -> Expr
f [Double]
xs = do
    let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs; e :: Expr
e = [Expr] -> Expr
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
    graph <- Int -> Expr -> Either MathError CompiledGraph
compileGraph Int
n Expr
e
    _ <- evaluateGraph graph xs
    ds <- mapM (`differentiate` e) [0 .. n - 1]
    mapM (\Expr
d -> ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient (Expr -> [Expr] -> Expr
forall a b. a -> b -> a
const Expr
d) [Double]
xs) ds

-- | Multiply the automatic Jacobian by a finite input direction; returns one value per output.
jvp :: ([Expr] -> [Expr]) -> [Double] -> [Double] -> Either MathError [Double]
jvp :: ([Expr] -> [Expr])
-> [Double] -> [Double] -> Either MathError [Double]
jvp [Expr] -> [Expr]
f [Double]
xs [Double]
ds = do
    [Double] -> [Double] -> Either MathError ()
forall a b. [a] -> [b] -> Either MathError ()
same [Double]
xs [Double]
ds
    j <- ([Expr] -> [Expr]) -> [Double] -> Either MathError [[Double]]
jacobian [Expr] -> [Expr]
f [Double]
xs
    finite ds
    let values = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ([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 b c a. (b -> c) -> (a -> b) -> a -> c
. (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]
ds) [[Double]]
j
    finite values
    Right values

-- | Multiply an output cotangent by the automatic Jacobian; returns one value per input.
vjp :: ([Expr] -> [Expr]) -> [Double] -> [Double] -> Either MathError [Double]
vjp :: ([Expr] -> [Expr])
-> [Double] -> [Double] -> Either MathError [Double]
vjp [Expr] -> [Expr]
f [Double]
xs [Double]
ds = do
    let ys :: [Expr]
ys = [Expr] -> [Expr]
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
    [Expr] -> [Double] -> Either MathError ()
forall a b. [a] -> [b] -> Either MathError ()
same [Expr]
ys [Double]
ds
    [Double] -> Either MathError ()
finite [Double]
ds
    ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient (Expr -> [Expr] -> Expr
forall a b. a -> b -> a
const ([Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Expr -> Expr) -> [Double] -> [Expr] -> [Expr]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
d Expr
y -> Double -> Expr
C Double
d Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
y) [Double]
ds [Expr]
ys))) [Double]
xs

-- | Multiply the symbolic Hessian by a checked finite input direction.
hvp :: ([Expr] -> Expr) -> [Double] -> [Double] -> Either MathError [Double]
hvp :: ([Expr] -> Expr)
-> [Double] -> [Double] -> Either MathError [Double]
hvp [Expr] -> Expr
f [Double]
xs [Double]
ds = do
    [Double] -> [Double] -> Either MathError ()
forall a b. [a] -> [b] -> Either MathError ()
same [Double]
xs [Double]
ds
    h <- ([Expr] -> Expr) -> [Double] -> Either MathError [[Double]]
hessian [Expr] -> Expr
f [Double]
xs
    finite ds
    let values = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ([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 b c a. (b -> c) -> (a -> b) -> a -> c
. (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]
ds) [[Double]]
h
    finite values
    Right values

{- | Explicit gradient accumulation owner; evaluation does not mutate it implicitly. Use the
accumulation API to retain/reset gradients.
-}
newtype GradientState = GradientState
    { GradientState -> [Double]
accumulatedGradient :: [Double]
    -- ^ Gradient components currently retained by explicit GradientState accumulation; reverse evaluation alone does not mutate it.
    }
    deriving (GradientState -> GradientState -> Bool
(GradientState -> GradientState -> Bool)
-> (GradientState -> GradientState -> Bool) -> Eq GradientState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: GradientState -> GradientState -> Bool
== :: GradientState -> GradientState -> Bool
$c/= :: GradientState -> GradientState -> Bool
/= :: GradientState -> GradientState -> Bool
Eq, Int -> GradientState -> ShowS
[GradientState] -> ShowS
GradientState -> String
(Int -> GradientState -> ShowS)
-> (GradientState -> String)
-> ([GradientState] -> ShowS)
-> Show GradientState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GradientState -> ShowS
showsPrec :: Int -> GradientState -> ShowS
$cshow :: GradientState -> String
show :: GradientState -> String
$cshowList :: [GradientState] -> ShowS
showList :: [GradientState] -> ShowS
Show)

-- | Create an explicit zero accumulation buffer for a nonnegative input count.
zeroGradients :: Int -> Either MathError GradientState
zeroGradients :: Int -> Either MathError GradientState
zeroGradients Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError GradientState
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative gradient size")
    | Bool
otherwise = GradientState -> Either MathError GradientState
forall a b. b -> Either a b
Right ([Double] -> GradientState
GradientState (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0))

-- | Add one same-length finite gradient without modifying the previous state.
accumulateGradients :: GradientState -> [Double] -> Either MathError GradientState
accumulateGradients :: GradientState -> [Double] -> Either MathError GradientState
accumulateGradients (GradientState [Double]
a) [Double]
b = do
    [Double] -> [Double] -> Either MathError ()
forall a b. [a] -> [b] -> Either MathError ()
same [Double]
a [Double]
b
    let c :: [Double]
c = (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
    [Double] -> Either MathError ()
finite [Double]
c
    GradientState -> Either MathError GradientState
forall a b. b -> Either a b
Right ([Double] -> GradientState
GradientState [Double]
c)

-- | Return a zero buffer of the original size; existing state remains reusable.
resetGradients :: GradientState -> GradientState
resetGradients :: GradientState -> GradientState
resetGradients (GradientState [Double]
a) = [Double] -> GradientState
GradientState ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a b. a -> b -> a
const Double
0) [Double]
a)

-- | Apply an explicit checked transformation; changing gradient length is rejected.
gradientHook ::
    ([Double] -> Either MathError [Double]) -> GradientState -> Either MathError GradientState
gradientHook :: ([Double] -> Either MathError [Double])
-> GradientState -> Either MathError GradientState
gradientHook [Double] -> Either MathError [Double]
f (GradientState [Double]
a) = do b <- [Double] -> Either MathError [Double]
f [Double]
a; same a b; finite b; Right (GradientState b)

-- | Clip Euclidean norm while preserving direction, including vectors whose unscaled norm exceeds Double range.
clipGradient :: Double -> [Double] -> Either MathError [Double]
clipGradient :: Double -> [Double] -> Either MathError [Double]
clipGradient Double
limit [Double]
xs
    | Bool -> Bool
not (Double -> Bool
isFinite Double
limit) Bool -> Bool -> Bool
|| Double
limit Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left MathError
InvalidTolerance
    | Bool
otherwise = 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
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]
xs)
            unitNorm :: Double
unitNorm = 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. Fractional a => a -> a -> a
/ Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs))
        [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right (if Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
|| Double
m Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
unitNorm then [Double]
xs else (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
unitNorm)) [Double]
xs)

{- | Returns maximum absolute discrepancy. The numerical check is deliberately
separate from automatic derivatives; h must be finite and strictly positive.
-}
gradientCheck :: Double -> ([Expr] -> Expr) -> [Double] -> Either MathError Double
gradientCheck :: Double -> ([Expr] -> Expr) -> [Double] -> Either MathError Double
gradientCheck Double
h [Expr] -> Expr
f [Double]
xs
    | Bool -> Bool
not (Double -> Bool
isFinite Double
h) Bool -> Bool -> Bool
|| Double
h Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = MathError -> Either MathError Double
forall a b. a -> Either a b
Left MathError
InvalidTolerance
    | Bool
otherwise = do
        g <- ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient [Expr] -> Expr
f [Double]
xs
        graph <- compileGraph (length xs) (f (map V [0 .. length xs - 1]))
        fd <-
            mapM
                ( \Int
i -> do
                    a <- CompiledGraph -> [Double] -> Either MathError Double
evaluateGraph CompiledGraph
graph (Int -> Double -> [Double] -> [Double]
forall {a}. Num a => Int -> a -> [a] -> [a]
bump Int
i Double
h [Double]
xs)
                    b <- evaluateGraph graph (bump i (-h) xs)
                    Right ((a - b) / (2 * h))
                )
                [0 .. length xs - 1]
        finite fd
        let discrepancies = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double -> Double
forall a. Num a => a -> a
abs (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b)) [Double]
g [Double]
fd
        finite discrepancies
        Right (maximum (0 : discrepancies))
  where
    bump :: Int -> a -> [a] -> [a]
bump Int
i a
h' [a]
ys = [if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then a
y a -> a -> a
forall a. Num a => a -> a -> a
+ a
h' else a
y | (Int
j, a
y) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [a]
ys]

-- | Differentiate a real objective of [real,imaginary] and return its Wirtinger derivative with respect to conjugate z.
complexGradient :: ([Expr] -> Expr) -> Complex Double -> Either MathError (Complex Double)
complexGradient :: ([Expr] -> Expr)
-> Complex Double -> Either MathError (Complex Double)
complexGradient [Expr] -> Expr
f (Double
x :+ Double
y) = do
    g <- ([Expr] -> Expr) -> [Double] -> Either MathError [Double]
gradient [Expr] -> Expr
f [Double
x, Double
y]
    case g of { [Double
a, Double
b] -> Complex Double -> Either MathError (Complex Double)
forall a b. b -> Either a b
Right (Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2); [Double]
_ -> MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"complex gradient arity") }

-- | Take a symbolic gradient-descent step at the supplied expressions; local substitution preserves outer derivatives across repeated steps.
differentiableStep :: Double -> ([Expr] -> Expr) -> [Expr] -> Either MathError [Expr]
differentiableStep :: Double -> ([Expr] -> Expr) -> [Expr] -> Either MathError [Expr]
differentiableStep Double
rate [Expr] -> Expr
f [Expr]
xs = do
    [Double] -> Either MathError ()
finite [Double
rate]
    let n :: Int
n = [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
xs
    ds <- (Int -> Either MathError Expr) -> [Int] -> Either MathError [Expr]
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 -> Int -> Expr -> Either MathError Expr
differentiate Int
i ([Expr] -> Expr
f ((Int -> Expr) -> [Int] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Expr
V [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    replacements <- mapM substitute ds
    Right (zipWith (\Expr
x Expr
d -> Expr
x Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Double -> Expr
C Double
rate Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
* Expr
d) xs replacements)
  where
    substitute :: Expr -> Either MathError Expr
substitute Expr
expression = case Expr
expression of
        C Double
x -> Expr -> Either MathError Expr
forall a b. b -> Either a b
Right (Double -> Expr
C Double
x)
        V Int
i ->
            if 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
< [Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
xs
                then Expr -> Either MathError Expr
forall a b. b -> Either a b
Right ([Expr]
xs [Expr] -> Int -> Expr
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                else MathError -> Either MathError Expr
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"differentiable step local variable index")
        Add Expr
a Expr
b -> Expr -> Expr -> Expr
Add (Expr -> Expr -> Expr)
-> Either MathError Expr -> Either MathError (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a Either MathError (Expr -> Expr)
-> Either MathError Expr -> Either MathError Expr
forall a b.
Either MathError (a -> b)
-> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Either MathError Expr
substitute Expr
b
        Mul Expr
a Expr
b -> Expr -> Expr -> Expr
Mul (Expr -> Expr -> Expr)
-> Either MathError Expr -> Either MathError (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a Either MathError (Expr -> Expr)
-> Either MathError Expr -> Either MathError Expr
forall a b.
Either MathError (a -> b)
-> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> Either MathError Expr
substitute Expr
b
        Neg Expr
a -> Expr -> Expr
Neg (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a
        Inv Expr
a -> Expr -> Expr
Inv (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a
        U Unary
op Expr
a -> Unary -> Expr -> Expr
U Unary
op (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a
        Stop Expr
a -> Expr -> Expr
Stop (Expr -> Expr) -> Either MathError Expr -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Either MathError Expr
substitute Expr
a
        Custom [Double] -> Either MathError (Double, [Double])
rule [Expr]
as -> ([Double] -> Either MathError (Double, [Double])) -> [Expr] -> Expr
Custom [Double] -> Either MathError (Double, [Double])
rule ([Expr] -> Expr)
-> Either MathError [Expr] -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr -> Either MathError Expr)
-> [Expr] -> Either MathError [Expr]
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 Expr -> Either MathError Expr
substitute [Expr]
as
        Checkpoint [Expr] -> Expr
recipe [Expr]
as -> ([Expr] -> Expr) -> [Expr] -> Expr
Checkpoint [Expr] -> Expr
recipe ([Expr] -> Expr)
-> Either MathError [Expr] -> Either MathError Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr -> Either MathError Expr)
-> [Expr] -> Either MathError [Expr]
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 Expr -> Either MathError Expr
substitute [Expr]
as

{- | Row-major symbolic tensors. Broadcasting is delegated to the checked tensor
shape engine. Matrix/contraction rules are ordinary scalar graph operations,
so reverse accumulation includes all repeated/broadcast input contributions.
-}
newtype ADTensor = ADTensor (T.Tensor Expr) deriving (Int -> ADTensor -> ShowS
[ADTensor] -> ShowS
ADTensor -> String
(Int -> ADTensor -> ShowS)
-> (ADTensor -> String) -> ([ADTensor] -> ShowS) -> Show ADTensor
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ADTensor -> ShowS
showsPrec :: Int -> ADTensor -> ShowS
$cshow :: ADTensor -> String
show :: ADTensor -> String
$cshowList :: [ADTensor] -> ShowS
showList :: [ADTensor] -> ShowS
Show)

-- | Construct a checked row-major symbolic tensor. Shape product must match expression count.
adTensor :: [Int] -> [Expr] -> Either MathError ADTensor
adTensor :: [Int] -> [Expr] -> Either MathError ADTensor
adTensor [Int]
shape [Expr]
xs = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> [Expr] -> Either MathError (Tensor Expr)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int]
shape [Expr]
xs

-- | Return logical row-major graph expressions, preserving shared input identifiers.
adValues :: ADTensor -> [Expr]
adValues :: ADTensor -> [Expr]
adValues (ADTensor Tensor Expr
t) = Tensor Expr -> [Expr]
forall a. Tensor a -> [a]
T.tensorValues Tensor Expr
t

-- | Return tensor dimensions, with [] denoting a scalar.
adShape :: ADTensor -> [Int]
adShape :: ADTensor -> [Int]
adShape (ADTensor Tensor Expr
t) = Tensor Expr -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Expr
t

-- | Broadcast through explicit repeated scalar graph leaves; reverse gradients sum over repetitions.
adBroadcast :: [Int] -> ADTensor -> Either MathError ADTensor
adBroadcast :: [Int] -> ADTensor -> Either MathError ADTensor
adBroadcast [Int]
s (ADTensor Tensor Expr
t) = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> Tensor Expr -> Either MathError (Tensor Expr)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
T.broadcastTo [Int]
s Tensor Expr
t

-- | Broadcast and combine symbolic tensors using the existing checked tensor shape contract.
adAdd, adMultiply :: ADTensor -> ADTensor -> Either MathError ADTensor
adAdd :: ADTensor -> ADTensor -> Either MathError ADTensor
adAdd (ADTensor Tensor Expr
a) (ADTensor Tensor Expr
b) = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tensor Expr -> Tensor Expr -> Either MathError (Tensor Expr)
forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
T.tensorAdd Tensor Expr
a Tensor Expr
b
adMultiply :: ADTensor -> ADTensor -> Either MathError ADTensor
adMultiply (ADTensor Tensor Expr
a) (ADTensor Tensor Expr
b) = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tensor Expr -> Tensor Expr -> Either MathError (Tensor Expr)
forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
T.tensorMultiply Tensor Expr
a Tensor Expr
b

-- | Symbolic batched matrix multiplication with checked contraction dimensions.
adMatmul :: ADTensor -> ADTensor -> Either MathError ADTensor
adMatmul :: ADTensor -> ADTensor -> Either MathError ADTensor
adMatmul (ADTensor Tensor Expr
a) (ADTensor Tensor Expr
b) = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tensor Expr -> Tensor Expr -> Either MathError (Tensor Expr)
forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
T.batchedMatmul Tensor Expr
a Tensor Expr
b

-- | Contract selected axes with matching sizes; uncontracted dimensions retain tensor order.
adContract :: [Int] -> [Int] -> ADTensor -> ADTensor -> Either MathError ADTensor
adContract :: [Int] -> [Int] -> ADTensor -> ADTensor -> Either MathError ADTensor
adContract [Int]
ax [Int]
bx (ADTensor Tensor Expr
a) (ADTensor Tensor Expr
b) = Tensor Expr -> ADTensor
ADTensor (Tensor Expr -> ADTensor)
-> Either MathError (Tensor Expr) -> Either MathError ADTensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
-> [Int]
-> Tensor Expr
-> Tensor Expr
-> Either MathError (Tensor Expr)
forall a.
Num a =>
[Int]
-> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
T.tensordot [Int]
ax [Int]
bx Tensor Expr
a Tensor Expr
b

-- | Reduce every tensor element to one symbolic scalar.
adSum :: ADTensor -> Expr
adSum :: ADTensor -> Expr
adSum = [Expr] -> Expr
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Expr] -> Expr) -> (ADTensor -> [Expr]) -> ADTensor -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADTensor -> [Expr]
adValues