ยง The smallest implementation of reverse mode AD (autograd) ever:
import qualified Data.Map.Strict as M
type Sym = String
type E a = M.Map Sym a
type F = Float
newtype Der = Der { under :: F } deriving(Show, Num)
infixl 7 !#
(!#) :: E a -> Sym -> a
(!#) = (M.!)
data Node =
Node { name :: Sym
, ins :: [Node]
, out :: E F -> F
, der :: (E F, E (Sym -> Der))
-> Sym -> Der
}
(!@) :: E a -> Node -> a
(!@) e node = e M.! (name node)
run_ :: (E F, E (Sym -> Der)) -> Node -> (E F, E (Sym -> Der))
run_ ein (Node name ins out der) =
let (e', ed') = foldl run_ ein ins
v = out e'
dv = der (e', ed')
in (M.insert name v e', M.insert name dv ed')
run :: E F -> Node -> (E F, E (Sym -> Der))
run e n = run_ (e, mempty) n
nconst :: Sym -> F -> Node
nconst n f = Node n [] (\_ -> f) (\_ _ -> 0)
nvar :: Sym -> Node
nvar n = Node n [] (!# n) (\_ n' -> if n == n' then 1 else 0)
nbinop :: (F -> F -> F)
-> (F -> Der -> F -> Der -> Der)
-> Sym
-> (Node, Node)
-> Node
nbinop f df n (in1, in2) =
Node { name = n
, ins = [in1, in2]
, out = \e -> f (e !# name in1) (e !# name in2)
, der = \(e, ed) n' ->
let (name1, name2) = (name in1, name in2)
(v1, v2) = (e !# name1, e !# name2)
(dv1, dv2) = (ed !# name1 $ n', ed !# name2 $ n')
in df v1 dv1 v2 dv2
}
nadd :: Sym -> (Node, Node) -> Node
nadd = nbinop (+) (\v dv v' dv' -> dv + dv')
nmul :: Sym -> (Node, Node) -> Node
nmul = nbinop (*) (\v (Der dv) v' (Der dv') -> Der $ (v*dv') + (v'*dv))
main :: IO ()
main = do
let x = nvar "x" :: Node
let y = nvar "y"
let xsq = nmul "xsq" (x, x)
let ten = nconst "10" 10
let xsq_plus_10 = nadd "xsq_plus_10" (xsq, ten)
let xsq_plus_10_plus_y = nadd "xsq_plus_10_plus_y" (xsq_plus_10, y)
let (e, de) = run (M.fromList $ [("x", 2.0), ("y", 3.0)]) xsq_plus_10_plus_y
putStrLn $ show e
putStrLn $ show $ de !@ xsq_plus_10_plus_y $ "x"
putStrLn $ show $ de !@ xsq_plus_10_plus_y $ "y"
Yeah, in ~80 lines of code, you can basically build an autograd engine. Isn't
haskell so rad?