
  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 | ||| Max sampler, always picks the element with the highest logit
 10 | public export
 11 | [pickMax] Monad m => MonadSample m where
 12 |   sample = toCostate $ \d => pure (argmax d.logits)
 13 | 
 14 | ||| Min sampler, always picks the element with the lowest logit
 15 | public export
 16 | [pickMin] Monad m => MonadSample m where
 17 |   sample = toCostate $ \d => pure (argmin d.logits)
 18 | 
 19 | ||| Sample an index of a cubical distribution
 20 | ||| Compute the cumulative distribution, draw uniformly, find the right bin 
 21 | public export
 22 | sampleCubical : {n : Nat} -> {name : AxisName} ->
 23 |   Tensor [name ~~> n] Double -> IO (Maybe (Fin n))
 24 | sampleCubical {n = Z} _ = pure Nothing
 25 | sampleCubical {n = S k} logits = do
 26 |   let cumSum = Utils.cumulativeSum (softargmaxImpl logits)
 27 |   r <- randomRIO (0.0, 1.0)
 28 |   pure $ findBin (#> cumSum) r
 29 | 
 30 | ||| Flattens a container distribution, sample the cubicla, map the index back
 31 | public export
 32 | MonadSample IO where
 33 |   sample @{ne} @{fold} = toCostate $ \d => do
 34 |     Just k <- sampleCubical (flatten mapToList d.logits)
 35 |       | Nothing => pure (GetInterface ne d.logits.extractShapeRank1) -- should never happen
 36 |     pure (mapToList.bwd d.logits.extractShapeRank1 k)
 37 |