0 | module Control.Monad.Distribution
  1 |
  2 | import Data.Vect
  3 | import Data.Fin
  4 | import Data.Bag
  5 |
  6 | import Data.Num
  7 | import public Data.Tensor
  8 | import Data.Container.Additive
  9 |
 10 | ||| Convex combination of a finite set of types, a point in a simplex △^(i-1)
 11 | ||| i=2 -> △¹ -> line segment
 12 | ||| i=3 -> △² -> triangle
 13 | ||| ...
 14 | ||| Probabilities are represented as logits, represented as a rank 1 tensor.
 15 | ||| Because the tensor has a name, and `Dist` is a thin wrapper around it,
 16 | ||| the name is exposed at the type level, allowing named operations to be
 17 | ||| extended to distributions
 18 | ||| TODO change description, as now `Dist` involves containers
 19 | ||| TODO, is Dist a quotient container?
 20 | public export
 21 | record Dist (a : Axis) where
 22 |   constructor MkDist
 23 |   ||| Probabilities are represented as logits
 24 |   logits : Tensor [a] Double
 25 |
 26 | ||| Logit representation of the uniform distribution
 27 | ||| Is defined for a container once the choice of shape is made
 28 | public export
 29 | uniform : {a : Axis} -> (s : a.cont.Shp) -> Dist a
 30 | uniform s = MkDist (fillAt (shapeToRank1Shape s) 0)
 31 |
 32 | ||| Logit representation of dirac delta at the position `p` of the shape `s`
 33 | ||| Note that `0` is the canonical choice, as softargmax subtracts the max
 34 | public export
 35 | diracDelta : {a : Axis} -> IsDecidable a.cont =>
 36 |   (s : a.cont.Shp) -> a.cont.Pos s -> Dist a
 37 | diracDelta s p = MkDist $ fillAt (shapeToRank1Shape s) minusInfinity // [p] .~ 0
 38 |
 39 | ||| Naperian axes have a single shape (`()`), so we do not need to supply it
 40 | namespace Naperian
 41 |   public export
 42 |   uniform : {a : Axis} -> IsNaperian a => Dist a
 43 |   uniform @{MkIsNaperian _ _} = uniform ()
 44 |
 45 |   public export
 46 |   diracDelta : {a : Axis} -> IsNaperian a => DecEq (Log a) => Log a -> Dist a
 47 |   diracDelta @{MkIsNaperian _ _} p = diracDelta () p
 48 |
 49 | namespace Cont
 50 |   ||| Container whose choice of a shape is a distribution over some choices, 
 51 |   ||| and a position is a specific choice made.
 52 |   public export
 53 |   Dist : Axis -> Cont
 54 |   Dist a = (d : Dist a) !> a.cont.Pos d.logits.extractShapeRank1
 55 |
 56 | public export
 57 | uniformLens : {a : Axis} -> a.cont =%> Cont.Dist a
 58 | uniformLens = !% \s => (uniform s ** id)
 59 |
 60 | ||| Container whose shapes are distributions (over positions over some chosen
 61 | ||| container shape), and positions are their gradients.
 62 | ||| Both are represented as logits
 63 | ||| If we were treating this as non-logit distributions then we'd have a
 64 | ||| one less dimension: both for the simplex in the forward pass and the
 65 | ||| gradients in the backwards one
 66 | ||| That is, the effective dimension of this space is n-1 (we can add a
 67 | ||| constant to all logits without changing the answer), and there's a
 68 | ||| direction in the gradient logit space that does not affect output
 69 | public export
 70 | Simplex : Axis -> AddCont
 71 | Simplex a = MkAddCont
 72 |   (Dist a) (\d =>
 73 |     (Tensor [a.name ~> (a.cont) `At` d.logits.extractShapeRank1] Double ** numIsMonoid))
 74 |
 75 | ||| Distributions are shown as probabilities (via softargmax), not as logits
 76 | public export
 77 | {a : Axis} ->
 78 | AllDisplay2D [a] Double =>
 79 | IsFoldable a.cont =>
 80 | AllAlgebra [a] Double =>
 81 | TensorCubEvidence [a] =>
 82 | Display2D (Dist a) where
 83 |   display2D (MkDist xs) = display2D (softargmaxImpl xs)
 84 |
 85 | public export
 86 | {a : Axis} ->
 87 | AllDisplay2D [a] Double =>
 88 | IsFoldable a.cont =>
 89 | AllAlgebra [a] Double =>
 90 | TensorCubEvidence [a] =>
 91 | Show (Dist a) where
 92 |   show = showViaDisplay2D
 93 |
 94 |
 95 | ||| A distribution over a container of branches, together with answers ready 
 96 | ||| when any branch is chosen, computed only when the branch is asked for
 97 | ||| TODO do we think of distr. on the fw pass as being part of Simplex or Nap?
 98 | public export
 99 | ProbabilisticChoice : (a : Axis) -> a.cont.filledWith AddCont -> AddCont
100 | ProbabilisticChoice a branches
101 |   = Simplex (a `At` (shapeExt branches)) >*< Section (index branches)
102 |
103 | ||| Each probabilistic choice can be turned into a sequence of producing a 
104 | ||| distribution, waiting for the response from the environment about which 
105 | ||| choice to make, then making that choice
106 | public export
107 | pendingChoice : {a : Axis} ->
108 |   (branches : a.cont.filledWith AddCont) ->
109 |   ProbabilisticChoice a branches =%+> Dist a >-+@ AddContDPair (index branches)
110 | pendingChoice branches =
111 |   let Choice : AddCont
112 |       Choice = Simplex (a `At` (shapeExt branches))
113 |   in (id {c = Choice} >*< graph) %+>> distribute
114 |   {c = Simplex (a `At` (shapeExt branches)),
115 |    g = AddContDPair (index branches)} (\d => !% \_ =>
116 |   (MkDist (MkT ((shapeExt branches <| const ()) <| \pp =>
117 |      index (GetT d.logits) (fst pp ** ()))) ** id))