
   0 | module Data.Autodiff.Model
   1 | 
   2 | import Data.Fin
   3 | import Data.Bag
   4 | import Data.ComMonoid
   5 | import Data.Container.Additive
   6 | import public Data.Para
   7 | import Data.Materialise
   8 | import public System.Random
   9 | 
  10 | {-------------------------------------------------------------------------------
  11 | Towards a typed analogue of nn.Module
  12 | -------------------------------------------------------------------------------}
  13 | 
  14 | -- todo need to rethink tihs syntax
  15 | export infixr 1 >>> -- sequential
  16 | export infixr 3 *** -- parallel
  17 | export infixr 3 &&& -- fan-out
  18 | 
  19 | ||| A Model is a differentiable parametric map which
  20 | ||| a) has its parameter container constant
  21 | ||| b) comes with its own initialisation
  22 | public export
  23 | record Model (a, b : AddCont) where
  24 |   constructor MkModel
  25 |   Params : Type
  26 |   {auto pMon : ComMonoid Params}
  27 |   init : IO Params
  28 |   run  : (a >*< Const Params @{pMon}) =%+> b
  29 | 
  30 | ||| Infix notation for the Model
  31 | public export
  32 | (-\->) : AddCont -> AddCont -> Type
  33 | a -\-> b = Model a b
  34 | 
  35 | public export
  36 | ParamCont : Model a b -> AddCont
  37 | ParamCont m = Const m.Params @{m.pMon}
  38 | 
  39 | namespace ParaConversion
  40 |   public export
  41 |   toPara : {0 a, b : AddCont} -> a -\-> b -> ParaAddLens a b
  42 |   toPara m = MkPara (ParamCont m) m.run
  43 |   
  44 |   public export
  45 |   fromPara : {0 a, b : AddCont} -> (f : ParaAddLens a b) ->
  46 |     (isConst : IsConst (Param f)) =>
  47 |     (init : IO (Param f).Shp) ->
  48 |     a -\-> b
  49 |   fromPara (MkPara _ f) {isConst = MkIsConst p @{mon}} init = MkModel p init f
  50 | 
  51 | 
  52 | ||| Replace a model's initialisation
  53 | public export
  54 | withInit : {0 a, b : AddCont} -> (m : a -\-> b) -> IO m.Params -> a -\-> b
  55 | withInit m i = MkModel m.Params @{m.pMon} i m.run
  56 | 
  57 | ||| Default parameter initialisation, uniform on (-1, 1)
  58 | public export
  59 | DefaultInit : Random p => Neg p => IO p
  60 | DefaultInit = randomRIO (-1, 1)
  61 | 
  62 | ||| A parameterless differentiable map
  63 | public export
  64 | trivialParam : {0 a, b : AddCont} -> a =%+> b -> a -\-> b
  65 | trivialParam f = MkModel Unit (pure ()) $
  66 |   !%+ \(x, ()) =>
  67 |     let (y ** k) = (%!+) f x
  68 |     in (y ** \y' => (k y', ()))
  69 | 
  70 | public export
  71 | id : {a : AddCont} -> a -\-> a
  72 | id = trivialParam id
  73 | 
  74 | ||| Sequential composition
  75 | public export
  76 | (>>>) : {0 a, b, c : AddCont} ->
  77 |   Materialise b.Shp => InterfaceOnPositions b Materialise =>
  78 |   a -\-> b ->
  79 |   b -\-> c ->
  80 |   a -\-> c
  81 | (MkModel p ip f) >>> (MkModel q iq g) = MkModel (p, q) [| (ip, iq) |] $
  82 |   (id >*< constPair)
  83 |     %+>> assocR {b=Const p, c=Const q}
  84 |     %+>> ((f %+>> materialiseCont) >*< id {c=Const q})
  85 |     %+>> g
  86 | 
  87 | ||| Parallel composition
  88 | public export
  89 | (***) : {a, b, c, d : AddCont} ->
  90 |   a -\-> c ->
  91 |   b -\-> d ->
  92 |   a >*< b -\-> c >*< d
  93 | (MkModel p ip f) *** (MkModel q iq g) = MkModel (p, q) [| (ip, iq) |] $
  94 |   (id {c=a>*<b} >*< constPair)
  95 |     %+>> swapMiddle {c3=Const p} {c4=Const q}
  96 |     %+>> (f >*< g)
  97 | 
  98 | ||| Fan-out
  99 | public export
 100 | (&&&) : {a : AddCont} -> {0 b, c : AddCont} ->
 101 |   a -\-> b ->
 102 |   a -\-> c ->
 103 |   a -\-> b >*< c
 104 | (MkModel p ip f) &&& (MkModel q iq g) = MkModel (p, q) [| (ip, iq) |] $
 105 |   (id >*< constPair)
 106 |     %+>> (copy >*< id {c=Const p >*< Const q})
 107 |     %+>> swapMiddle {c3=Const p} {c4=Const q}
 108 |     %+>> (f >*< g)
 109 | 
 110 | ||| Fan-out of models
 111 | public export
 112 | dfunFinite : {a : AddCont} -> {n : Nat} -> {0 f : Fin n -> AddCont} ->
 113 |   ((i : Fin n) -> a -\-> f i) -> a -\-> AddContDFunFinite f
 114 | dfunFinite {n = 0} _ = trivialParam terminal
 115 | dfunFinite {n = S k} ms = ms 0 &&& dfunFinite {f = f . FS} (\i => ms (FS i))
 116 | 
 117 | ||| Only evaluates the head if the index matches it
 118 | public export
 119 | lazyCons : {a : AddCont} ->
 120 |   {0 b : AddCont} -> {0 k : Nat} -> {0 bs : Vect k AddCont} ->
 121 |   a -\-> b -> a -\-> (Vect k >-+@ Coproduct bs) ->
 122 |   a -\-> (Vect (S k) >-+@ Coproduct (b :: bs))
 123 | lazyCons (MkModel p @{pm} ip f) (MkModel q @{qm} iq g) = MkModel
 124 |   (p, q) [| (ip, iq) |] $
 125 |   !%+ \(x, (px, qx)) =>
 126 |     let hd : Lazy (t : b.Shp ** b.PosSet t -> (a >*< Const p).PosSet (x, px))
 127 |         hd = (%!+) f (x, px)
 128 |         rest = (%!+) g (x, qx)
 129 |     in (() <| (\case
 130 |             FZ => (FZ ** fst hd)
 131 |             FS j => (FS (fst (index (fst rest) j)) ** snd (index (fst rest) j)))
 132 |         ** fromGenerators {y = (a >*< Const (p, q)).Pos (x, (px, qx))}
 133 |              (\(i ** gr) => case i of
 134 |                 FZ => let (x', p') = snd hd gr
 135 |                       in (x', (p', (Const q).Zero qx))
 136 |                 FS j => let (x', q') = snd rest (MkBag [(j ** gr)])
 137 |                         in (x', ((Const p @{pm}).Zero px, q'))))
 138 | 
 139 | ||| Branch models under the choice effect: only the branch the environment asks
 140 | ||| for runs. The type of `postcomposeLens (dfunFinite ms) graph`, without its work
 141 | public export
 142 | lazyBranches : {a : AddCont} -> {n : Nat} -> {0 branches : Vect n AddCont} ->
 143 |   ((i : Fin n) -> a -\-> index i branches) ->
 144 |   a -\-> (Vect n >-+@ Coproduct branches)
 145 | lazyBranches {n = 0} {branches = []} _ = trivialParam $
 146 |   !%+ \x => (() <| (\i => absurd i) ** fromGenerators (\(i ** _) => absurd i))
 147 | lazyBranches {n = S k} {branches = b :: bs} ms
 148 |   = lazyCons (ms 0) (lazyBranches (\i => ms (FS i)))
 149 | 
 150 | ||| Act on the first component
 151 | public export
 152 | mapFst : {a, c : AddCont} ->
 153 |   a -\-> b ->
 154 |   a >*< c -\-> b >*< c
 155 | mapFst m = MkModel m.Params @{m.pMon} m.init $
 156 |   assocL {c=ParamCont m}
 157 |     %+>> (id >*< swap {a=c} {b=ParamCont m})
 158 |     %+>> assocR {a} {b=ParamCont m}
 159 |     %+>> (m.run >*< id)
 160 | 
 161 | ||| Iterate a model `n` times
 162 | public export
 163 | nTimes : {a : AddCont} ->
 164 |   Materialise a.Shp =>
 165 |   InterfaceOnPositions a Materialise => 
 166 |   Nat -> a -\-> a -> a -\-> a
 167 | nTimes 0 m = id
 168 | nTimes 1 m = m
 169 | nTimes (S k) m = m >>> nTimes k m
 170 | 
 171 | public export
 172 | postcomposeLens : {0 a, b, c : AddCont} -> a -\-> b -> b =%+> c -> a -\-> c
 173 | postcomposeLens m g = MkModel m.Params @{m.pMon} m.init (m.run %+>> g)
 174 | 
 175 | ||| Pre-compose a parameterless lens onto a model's input
 176 | public export
 177 | precomposeLens : {0 a, b, c : AddCont} -> a =%+> b -> b -\-> c -> a -\-> c
 178 | precomposeLens g m
 179 |   = MkModel m.Params @{m.pMon} m.init ((g >*< id {c = ParamCont m}) %+>> m.run)
 180 | 
 181 | ||| A custom function without parameters
 182 | public export
 183 | prim : {0 a : AddCont} -> {b : AddCont} ->
 184 |   ((x : a.Shp) -> (y : b.Shp ** (b.PosSet y -> a.PosSet x))) ->
 185 |   a -\-> b
 186 | prim f = trivialParam (!%+ f)
 187 | 
 188 | ||| A custom differentiable operation
 189 | public export
 190 | customOp : {s, t : Type} -> ComMonoid s => ComMonoid t =>
 191 |   (fwd : s -> t) -> (vjp : s -> t -> s) ->
 192 |   Const s -\-> Const t
 193 | customOp fwd vjp = prim (\x => (fwd x ** vjp x))
 194 | 
 195 | ||| A custom parametric layer
 196 | public export
 197 | layer : {0 a : AddCont} -> {b : AddCont} ->
 198 |   (p : Type) -> ComMonoid p =>
 199 |   (initP : IO p) ->
 200 |   ((x : a.Shp) -> (param : p) -> (y : b.Shp ** (b.PosSet y -> (a.PosSet x, p)))) ->
 201 |   a -\-> b
 202 | layer p initP f = MkModel p initP $
 203 |   !%+ \(x, ps) => f x ps
 204 | 
 205 | ||| Run a model at an input and a parameter
 206 | public export
 207 | runAt : {0 a, b : AddCont} -> (m : a -\-> b) ->
 208 |   (x : a.Shp) -> (p : m.Params) ->
 209 |   (y : b.Shp ** (b.PosSet y -> (a.PosSet x, m.Params)))
 210 | runAt m x p = (%!+) m.run (x, p)
 211 | 
 212 | ||| Forward pass
 213 | public export
 214 | (.fwd) : {0 a, b : AddCont} -> (m : a -\-> b) ->
 215 |   (x : a.Shp) -> (p : m.Params) -> b.Shp
 216 | (.fwd) m x p = fst (runAt m x p)
 217 | 
 218 | ||| Backward pass
 219 | public export
 220 | (.bwd) : {0 a, b : AddCont} -> (m : a -\-> b) ->
 221 |   (x : a.Shp) -> (p : m.Params) ->
 222 |   b.PosSet (m.fwd x p) -> (a.PosSet x, m.Params)
 223 | (.bwd) m x p = snd (runAt m x p)