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 |
111 | ||| Initialise every model of a family, into the tuple of their parameters
112 | public export
113 | initAll : {n : Nat} -> {0 a : AddCont} -> {0 f : Fin n -> AddCont} ->
114 |   (ms : (i : Fin n) -> a -\-> f i) ->
115 |   IO (Product (\i => ParamCont (ms i))).Shp
116 | initAll {n = 0} _ = pure ()
117 | initAll {n = S k} ms = [| ((ms 0).init, initAll (\i => ms (FS i))) |]
118 |
119 | ||| A family of models into the section of their codomains: the n-ary lazy
120 | ||| fan-out, a read running only the asked branch, with every branch's
121 | ||| parameters as one tuple
122 | public export
123 | fanOutModel : {a : AddCont} -> {n : Nat} -> {0 f : Fin n -> AddCont} ->
124 |   (ms : (i : Fin n) -> a -\-> f i) -> a -\-> Section f
125 | fanOutModel ms = MkModel
126 |   (Product (\i => ParamCont (ms i))).Shp
127 |   @{finiteShpMon {f = \i => ParamCont (ms i)} (\i => (ms i).pMon)}
128 |   (initAll ms)
129 |   ((id >*< constFinite {ps = \i => (ms i).Params} (\i => (ms i).pMon))
130 |     %+>> fanOut (\i => (id >*< projFinite {f = \i => ParamCont (ms i)} i) %+>> (ms i).run))
131 |
132 | ||| Iterate a model `n` times
133 | public export
134 | nTimes : {a : AddCont} ->
135 |   Materialise a.Shp =>
136 |   InterfaceOnPositions a Materialise => 
137 |   Nat -> a -\-> a -> a -\-> a
138 | nTimes 0 m = id
139 | nTimes 1 m = m
140 | nTimes (S k) m = m >>> nTimes k m
141 |
142 | public export
143 | postcomposeLens : {0 a, b, c : AddCont} -> a -\-> b -> b =%+> c -> a -\-> c
144 | postcomposeLens m g = MkModel m.Params @{m.pMon} m.init (m.run %+>> g)
145 |
146 | ||| Pre-compose a parameterless lens onto a model's input
147 | public export
148 | precomposeLens : {0 a, b, c : AddCont} -> a =%+> b -> b -\-> c -> a -\-> c
149 | precomposeLens g m
150 |   = MkModel m.Params @{m.pMon} m.init ((g >*< id {c = ParamCont m}) %+>> m.run)
151 |
152 | ||| A custom function without parameters
153 | public export
154 | prim : {0 a : AddCont} -> {b : AddCont} ->
155 |   ((x : a.Shp) -> (y : b.Shp ** (b.PosSet y -> a.PosSet x))) ->
156 |   a -\-> b
157 | prim f = trivialParam (!%+ f)
158 |
159 | ||| A custom differentiable operation
160 | public export
161 | customOp : {s, t : Type} -> ComMonoid s => ComMonoid t =>
162 |   (fwd : s -> t) -> (vjp : s -> t -> s) ->
163 |   Const s -\-> Const t
164 | customOp fwd vjp = prim (\x => (fwd x ** vjp x))
165 |
166 | ||| A custom parametric layer
167 | public export
168 | layer : {0 a : AddCont} -> {b : AddCont} ->
169 |   (p : Type) -> ComMonoid p =>
170 |   (initP : IO p) ->
171 |   ((x : a.Shp) -> (param : p) -> (y : b.Shp ** (b.PosSet y -> (a.PosSet x, p)))) ->
172 |   a -\-> b
173 | layer p initP f = MkModel p initP $
174 |   !%+ \(x, ps) => f x ps
175 |
176 | ||| Run a model at an input and a parameter
177 | public export
178 | runAt : {0 a, b : AddCont} -> (m : a -\-> b) ->
179 |   (x : a.Shp) -> (p : m.Params) ->
180 |   (y : b.Shp ** (b.PosSet y -> (a.PosSet x, m.Params)))
181 | runAt m x p = (%!+) m.run (x, p)
182 |
183 | ||| Forward pass
184 | public export
185 | (.fwd) : {0 a, b : AddCont} -> (m : a -\-> b) ->
186 |   (x : a.Shp) -> (p : m.Params) -> b.Shp
187 | (.fwd) m x p = fst (runAt m x p)
188 |
189 | ||| Backward pass
190 | public export
191 | (.bwd) : {0 a, b : AddCont} -> (m : a -\-> b) ->
192 |   (x : a.Shp) -> (p : m.Params) ->
193 |   b.PosSet (m.fwd x p) -> (a.PosSet x, m.Params)
194 | (.bwd) m x p = snd (runAt m x p)