0 | module Data.Autodiff.Ops
  1 |
  2 | import public System.Random
  3 |
  4 | import Data.Tensor
  5 | import Data.Tensor.Utils
  6 | import Data.Container.Additive
  7 | import Data.Container.Additive.Quantifiers
  8 | import Data.Para
  9 | import Data.Autodiff.Model
 10 | import Control.Monad.Distribution
 11 | import Data.ComMonoid
 12 | import Data.Materialise
 13 |
 14 | import Misc
 15 |
 16 | %hide Data.Container.Base.Morphism.Definition.DependentLenses.(=%>)
 17 |
 18 | {-------------------------------------------------------------------------------
 19 | This file contains the derivatives of various useful building blocks
 20 |
 21 | Eventually will be combined with functionality which functorially assigns these
 22 | to any forward pass
 23 |
 24 | -------------------------------------------------------------------------------}
 25 |
 26 | public export
 27 | mulModel : {t : Type} -> Neg t => Random t =>
 28 |   Const t -\-> Const t
 29 | mulModel = fromPara (binaryOpToPara {p = Const t} mul) DefaultInit
 30 |
 31 | public export
 32 | addModel : {t : Type} -> Neg t => Random t =>
 33 |   Const t -\-> Const t
 34 | addModel = fromPara (binaryOpToPara {p = Const t} sum) DefaultInit
 35 |
 36 | public export
 37 | scalarAffine : {t : Type} -> Neg t => Random t => Materialise t =>
 38 |   Const t -\-> Const t
 39 | scalarAffine = mulModel >>> addModel
 40 |
 41 | ||| Apply a scalar lens elementwise across a tensor
 42 | public export
 43 | parallelTensor : {a, b : Type} -> Num a => Num b =>
 44 |   {shape : TensorShape rank} ->
 45 |   AllC TensorMonoid shape => AllC IsConcrete shape =>
 46 |   (f : Const a =%+> Const b) ->
 47 |   Const (Tensor shape a) =%+> Const (Tensor shape b)
 48 | parallelTensor f = !%+ \x =>
 49 |   let outs : Tensor shape (y : b ** (b -> a))
 50 |       outs = materialise ((%!+) f <$> x)
 51 |   in (materialise (fst <$> outs) ** \ys' => materialise [| snd outs ys' |])
 52 |
 53 | ||| `Tensor shape` applied to a model: one copy per entry, each with its own parameter
 54 | public export
 55 | parallelTensorModel : {a, b : Type} -> Num a => Num b =>
 56 |   {shape : TensorShape rank} -> AllC TensorMonoid shape => AllC IsConcrete shape =>
 57 |   Traversable (Tensor shape) =>
 58 |   (m : Const a -\-> Const b) -> Const (Tensor shape a) -\-> Const (Tensor shape b)
 59 | parallelTensorModel m = MkModel (Tensor shape m.Params) @{tensorComMonoid m.pMon}
 60 |   (sequence (pure m.init)) $
 61 |   !%+ \(x, ps) =>
 62 |     let outs : Tensor shape (y : b ** b -> (a, m.Params))
 63 |         outs = materialise ([| (x, ps) |] <&> (%!+) m.run)
 64 |     in (materialise (fst <$> outs) ** \ys' =>
 65 |       let grads : Tensor shape (a, m.Params)
 66 |           grads = materialise [| snd outs ys' |]
 67 |       in (fst <$> grads, snd <$> grads))
 68 |
 69 | public export
 70 | copyN : {a : Type} -> Num a => {n : Nat} -> {axisName : AxisName} ->
 71 |   Const a =%+> Const (Tensor [axisName ~~> n] a)
 72 | copyN = !%+ \x => (pure x ** reduce)
 73 |
 74 | public export
 75 | sameFromTensorN : {a, b : Type} -> Num a => Num b => {n : Nat} ->
 76 |   {axisName : AxisName} ->
 77 |   Traversable (Tensor [axisName ~~> n]) =>
 78 |   (m : Const a -\-> Const b) -> Const a -\-> Const (Tensor [axisName ~~> n] b)
 79 | sameFromTensorN m = precomposeLens copyN (parallelTensorModel m)
 80 |
 81 | ||| Dual to `copyN`
 82 | public export
 83 | sumAxis : {n : Axis} -> IsCubical n => Num a =>
 84 |   TensorMonoid n.cont =>
 85 |   Const (Tensor [n] a) =%+> Const (Tensor [] a)
 86 | sumAxis @{MkIsCubical _ n} = !%+ \t => (># reduce t ** \a' => fill (#> a'))
 87 |
 88 | ||| Divide by a constant, entrywise in both directions
 89 | public export
 90 | divBy : {a : Type} -> Num a => Fractional a =>
 91 |   (d : a) ->
 92 |   Const (Tensor [] a) =%+> Const (Tensor [] a)
 93 | divBy d = !%+ \x => (x <&> (/ d) ** \x' => x' <&> (/ d))
 94 |
 95 | public export
 96 | meanSquaredDifference : IsCubical n => TensorMonoid n.cont =>
 97 |   {a : Type} -> Num a => Neg a => Fractional a => Cast Nat a =>
 98 |   Const (Tensor [n] a) >*< Const (Tensor [n] a) =%+> Const (Tensor [] a)
 99 | meanSquaredDifference @{MkIsCubical _ n}
100 |   = SquaredDifference %+>> sumAxis %+>> divBy (cast n)
101 |
102 | -- Activations
103 |
104 | ||| Recovers `ReLU` when `alpha=0`
105 | ||| Cannot be written as a composition of scaling and `ReLU`
106 | public export
107 | leakyReLU : {a : Type} -> Num a => Ord a =>
108 |   (alpha : a) ->
109 |   Const a =%+> Const a
110 | leakyReLU alpha = !%+ \x =>
111 |   (if x > 0 then x else alpha * x ** \x' => if x > 0 then x' else alpha * x')
112 |
113 | public export
114 | leakyReLUModel : {a : Type} -> Num a => Ord a =>
115 |   (alpha : a) ->
116 |   {shape : TensorShape rank} -> AllC TensorMonoid shape => AllC IsConcrete shape =>
117 |   Const (Tensor shape a) -\-> Const (Tensor shape a)
118 | leakyReLUModel alpha = trivialParam (parallelTensor (leakyReLU alpha))
119 |
120 | public export
121 | reluModel : {a : Type} -> Num a => Ord a =>
122 |   {shape : TensorShape rank} -> AllC TensorMonoid shape => AllC IsConcrete shape =>
123 |   Const (Tensor shape a) -\-> Const (Tensor shape a)
124 | reluModel = leakyReLUModel 0
125 |
126 | -- Distributions
127 |
128 | ||| Interpret a vector as logits of a distribution. The backward pass is
129 | ||| identity: gradients are computed in the sme way
130 | public export
131 | fromLogits : {0 name : AxisName} -> {0 n : Nat} ->
132 |   Const (Tensor [name ~~> n] Double) =%+> Simplex name n
133 | fromLogits = !%+ \xs => (MkDist xs ** id)
134 |
135 | public export
136 | fromLogitsModel : {0 name : AxisName} -> {0 n : Nat} ->
137 |   Const (Tensor [name ~~> n] Double) -\-> Simplex name n
138 | fromLogitsModel = trivialParam fromLogits
139 |