
  0 | module Control.Monad.Sample.Instances
  1 | 
  2 | import Control.Monad.Identity
  3 | import System.Random
  4 | 
  5 | import Data.Tensor
  6 | import Control.Monad.Distribution
  7 | import Control.Monad.Sample.Definition
  8 | 
  9 | ||| Trivial sampler, always picks the first element
 10 | public export
 11 | [pickFirst] MonadSample Identity where
 12 |   sample {i = (S k)} (MkDist xs) = Id FZ
 13 | 
 14 | ||| Max sampler, always picks the element with the highest logit
 15 | public export
 16 | [pickMax] MonadSample Identity where
 17 |   sample {i = (S k)} d = Id (argmax d.logits)
 18 | 
 19 | ||| Min sampler, always picks the element with the lowest logit
 20 | public export
 21 | [pickMin] MonadSample Identity where
 22 |   sample {i = (S k)} d = Id (argmin d.logits)
 23 | 
 24 | 
 25 | ||| Computes the cumulative distribution, samples randomly, finds the right bin
 26 | public export
 27 | MonadSample IO where
 28 |   sample {i = S j} (MkDist xs) = do
 29 |     let dist = softargmaxImpl xs
 30 |         cumSum = cumulativeSum dist
 31 |     r <- randomRIO (0.0, 1.0)
 32 |     case findBin (#> cumSum) r of
 33 |       Nothing => pure FZ -- should never happen!
 34 |       Just i => pure i
 35 | 
 36 | testIO : IO ()
 37 | testIO = do
 38 |   let logits : Dist "coin" 2
 39 |       logits = MkDist (># [-(1.099), 1.099]) -- this produces the dist [0.1, 0.9]
 40 |   is <- sequence (replicate 1000 (sample logits))
 41 |   -- printLn is
 42 |   printLn (count (== 0) is) -- should be ~100
 43 |   printLn (count (== 1) is) -- should be ~900
 44 | 
 45 | public export
 46 | testDirac : IO ()
 47 | testDirac = do
 48 |   let index = 4
 49 |   let logits = diracDelta {name="dirac"} {i=10} index
 50 |   inds <- sequence (replicate 1000 (sample logits))
 51 |   printLn (take 10 inds)
 52 |   printLn (count (== index) inds)