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, is Dist a quotient container?
19 | public export
20 | record Dist (name : AxisName) (i : Nat) where
21 |   constructor MkDist
22 |   ||| Probabilities are represented as logits
23 |   logits : Tensor [name ~~> i] Double
24 |
25 | ||| Logit representation of the uniform distribution
26 | public export
27 | uniform : {name : AxisName} -> {i : Nat} ->
28 |   (isSucc : IsSucc i) => Dist name i
29 | uniform = MkDist (fill 0)
30 |
31 | ||| Logit representation of dirac delta
32 | ||| Note that `0` is the canonical choice, as softargmax subtracts the max
33 | public export
34 | diracDelta : {name : AxisName} -> {i : Nat} ->
35 |   IsSucc i =>
36 |   (j : Fin i) -> Dist name i
37 | diracDelta @{ItIsSucc {n}} j
38 |   = MkDist (># (insertAt j 0 (replicate n minusInfinity)))
39 |
40 | namespace Cont
41 |   ||| Container whose shape represents a distribution over `n` choices, and
42 |   ||| whose position represents the choice made.
43 |   public export
44 |   Dist : AxisName -> Nat -> Cont
45 |   Dist name n = Const2 (Dist name n) (Fin n)
46 |
47 | ||| Container whose shapes are distributions, positions their gradients.
48 | ||| Both are represented as logits
49 | ||| If we were treating this as non-logit distributions then we'd have a
50 | ||| one less dimension: both for the simplex in the forward pass and the
51 | ||| gradients in the backwards one
52 | ||| That is, the effective dimension of this space is n-1 (we can add a
53 | ||| constant to all logits without changing the answer), and there's a
54 | ||| direction in the gradient logit space that does not affect output
55 | public export
56 | Simplex : AxisName -> Nat -> AddCont
57 | Simplex name n = Const2 (Dist name n)
58 |   (Tensor [name ~~> n] Double ** numIsMonoid)
59 |
60 | ||| Distributions are shown as probabilities (via softargmax), not as logits
61 | public export
62 | {axisName : AxisName} -> {i : Nat} -> Show (Dist axisName i) where
63 |   show (MkDist xs) = show (softargmaxImpl xs)
64 |
65 | ||| A distribution over `n` branches together with, contingent on a choice made
66 | ||| by the environment, the chosen branch's content
67 | ||| TODO do we think of distr. on the fw pass as being part of Simplex or Nap?
68 | public export
69 | ProbabilisticChoice : {n : Nat} -> (distName : AxisName) -> 
70 |   (branches : Vect n AddCont) -> AddCont
71 | ProbabilisticChoice distName branches
72 |   = Simplex distName n >*< (Vect n >-+@ Coproduct branches)
73 |
74 | ||| The choice made: a distribution over the branches and the chosen branch's content
75 | public export
76 | ChoiceMade : {n : Nat} -> (distName : AxisName) ->
77 |   (branches : Vect n AddCont) -> AddCont
78 | ChoiceMade distName branches = Simplex distName n >*< Coproduct branches
79 |
80 | ||| Resolve probabilistic choice through the ground truth label. The labelled 
81 | ||| branch is selected, and its gradient goes back as a singleton bag
82 | ||| TODO do we need the right component of the codomain?
83 | public export
84 | resolveByLabel : {distName : AxisName} ->
85 |   {branches : Vect n AddCont} ->
86 |   ProbabilisticChoice distName branches >*< ChoiceMade distName branches
87 |     =%+> ChoiceMade distName branches >*< ChoiceMade distName branches
88 | resolveByLabel = !%+ \((dist, ex), y@(distTrue, (iTrue ** _))) =>
89 |   (((dist, index ex iTrue), y) ** \((d', g'), yGrad) =>
90 |       ((d', MkBag [(iTrue ** g')]), yGrad))
91 |