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