diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,11 @@
+# Unreleased
+
+- No unreleased changes so far.
+
+# 0.1.1.0 (2020-04-08)
+
+- New exported function: `Control.Monad.Bayes.Class` now exports `discrete`.
+
 # 0.1.0.0 (2020-02-17)
 
 Initial release.
diff --git a/LICENSE b/LICENSE
deleted file mode 100644
--- a/LICENSE
+++ /dev/null
@@ -1,22 +0,0 @@
-The MIT License (MIT)
-
-Copyright (c) 2015-2020 Adam Scibior
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
diff --git a/LICENSE.md b/LICENSE.md
new file mode 100644
--- /dev/null
+++ b/LICENSE.md
@@ -0,0 +1,22 @@
+The MIT License (MIT)
+
+Copyright (c) 2015-2020 Adam Scibior
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
diff --git a/Setup.hs b/Setup.hs
--- a/Setup.hs
+++ b/Setup.hs
@@ -1,2 +1,3 @@
 import Distribution.Simple
+
 main = defaultMain
diff --git a/benchmark/SSM.hs b/benchmark/SSM.hs
--- a/benchmark/SSM.hs
+++ b/benchmark/SSM.hs
@@ -1,15 +1,13 @@
 module Main where
 
-import Control.Monad.IO.Class
-
-import Control.Monad.Bayes.Sampler
-import Control.Monad.Bayes.Weighted
-import Control.Monad.Bayes.Population
-import Control.Monad.Bayes.Inference.SMC
-import Control.Monad.Bayes.Inference.RMSMC
 import Control.Monad.Bayes.Inference.PMMH as PMMH
+import Control.Monad.Bayes.Inference.RMSMC
+import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Inference.SMC2 as SMC2
-
+import Control.Monad.Bayes.Population
+import Control.Monad.Bayes.Sampler
+import Control.Monad.Bayes.Weighted
+import Control.Monad.IO.Class
 import NonlinearSSM
 
 main :: IO ()
@@ -17,19 +15,15 @@
   let t = 5
   dat <- generateData t
   let ys = map snd dat
-
   liftIO $ print "SMC"
   smcRes <- runPopulation $ smcMultinomial t 10 (param >>= model ys)
   liftIO $ print $ show smcRes
-
   liftIO $ print "RM-SMC"
   smcrmRes <- runPopulation $ rmsmcLocal t 10 10 (param >>= model ys)
   liftIO $ print $ show smcrmRes
-
   liftIO $ print "PMMH"
   pmmhRes <- prior $ pmmh 2 t 3 param (model ys)
   liftIO $ print $ show pmmhRes
-
   liftIO $ print "SMC2"
   smc2Res <- runPopulation $ smc2 t 3 2 1 param (model ys)
   liftIO $ print $ show smc2Res
diff --git a/benchmark/Single.hs b/benchmark/Single.hs
--- a/benchmark/Single.hs
+++ b/benchmark/Single.hs
@@ -1,44 +1,42 @@
-import System.Random.MWC (createSystemRandom, GenIO)
-import Data.Time
-import Options.Applicative
-import Data.Semigroup ((<>))
-
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Sampler
-import Control.Monad.Bayes.Weighted
-import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Inference.RMSMC
+import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Population
-import Control.Monad.Bayes.Sequential
+import Control.Monad.Bayes.Sampler
 import Control.Monad.Bayes.Traced
-
+import Control.Monad.Bayes.Weighted
+import Data.Time
 import qualified HMM
-import qualified LogReg
 import qualified LDA
+import qualified LogReg
+import Options.Applicative
+import System.Random.MWC (createSystemRandom)
 
 data Model = LR Int | HMM Int | LDA (Int, Int)
-  deriving(Show, Read)
+  deriving (Show, Read)
+
 parseModel :: String -> Maybe Model
 parseModel s =
   case s of
-    'L':'R':n -> Just $ LR (read n)
-    'H':'M':'M':n -> Just $ HMM (read n)
-    'L':'D':'A':n -> Just $ LDA (5, read n)
+    'L' : 'R' : n -> Just $ LR (read n)
+    'H' : 'M' : 'M' : n -> Just $ HMM (read n)
+    'L' : 'D' : 'A' : n -> Just $ LDA (5, read n)
     _ -> Nothing
 
-getModel :: MonadInfer m => Model ->  (Int, m String)
-getModel model = (size model, program model) where
-  size (LR n) = n
-  size (HMM n) = n
-  size (LDA (d,w)) = d*w
-  synthesize :: SamplerST a -> (a -> b) -> b
-  synthesize dataGen prog = prog (sampleSTfixed dataGen)
-  program (LR n) = show <$> synthesize (LogReg.syntheticData n) LogReg.logisticRegression
-  program (HMM n) = show <$> synthesize (HMM.syntheticData n) HMM.hmm
-  program (LDA (d,w)) = show <$> synthesize (LDA.syntheticData d w) LDA.lda
+getModel :: MonadInfer m => Model -> (Int, m String)
+getModel model = (size model, program model)
+  where
+    size (LR n) = n
+    size (HMM n) = n
+    size (LDA (d, w)) = d * w
+    synthesize :: SamplerST a -> (a -> b) -> b
+    synthesize dataGen prog = prog (sampleSTfixed dataGen)
+    program (LR n) = show <$> synthesize (LogReg.syntheticData n) LogReg.logisticRegression
+    program (HMM n) = show <$> synthesize (HMM.syntheticData n) HMM.hmm
+    program (LDA (d, w)) = show <$> synthesize (LDA.syntheticData d w) LDA.lda
 
 data Alg = SMC | MH | RMSMC
-  deriving(Read,Show)
+  deriving (Read, Show)
 
 runAlg :: Model -> Alg -> SamplerIO String
 runAlg model alg =
@@ -46,17 +44,16 @@
     SMC ->
       let n = 100
           (k, m) = getModel model
-      in show <$> (runPopulation $ smcSystematic k n m)
-    MH  ->
+       in show <$> runPopulation (smcSystematic k n m)
+    MH ->
       let t = 100
           (_, m) = getModel model
-      in show <$> (prior $ mh t m)
+       in show <$> prior (mh t m)
     RMSMC ->
       let n = 10
           t = 1
           (k, m) = getModel model
-      in show <$> (runPopulation $ rmsmcBasic k n t m)
-
+       in show <$> runPopulation (rmsmcBasic k n t m)
 
 infer :: Model -> Alg -> IO ()
 infer model alg = do
@@ -65,23 +62,27 @@
   print x
 
 opts :: ParserInfo (Model, Alg)
-opts = flip info fullDesc $ liftA2 (,) model alg where
-  model = option (maybeReader parseModel)
-          ( long "model"
-          <> short 'm'
-          <> help "Model")
-  alg = option auto
-          ( long "alg"
-         <> short 'a'
-         <> help "Inference algorithm")
+opts = flip info fullDesc $ liftA2 (,) model alg
+  where
+    model =
+      option
+        (maybeReader parseModel)
+        ( long "model"
+            <> short 'm'
+            <> help "Model"
+        )
+    alg =
+      option
+        auto
+        ( long "alg"
+            <> short 'a'
+            <> help "Inference algorithm"
+        )
 
 main :: IO ()
 main = do
   (model, alg) <- execParser opts
-
   startTime <- getCurrentTime
-
   infer model alg
-
   endTime <- getCurrentTime
   print (diffUTCTime endTime startTime)
diff --git a/benchmark/Speed.hs b/benchmark/Speed.hs
--- a/benchmark/Speed.hs
+++ b/benchmark/Speed.hs
@@ -1,26 +1,21 @@
-import Criterion.Main
-import Criterion.Types
-import System.Random.MWC (createSystemRandom, GenIO)
-import Control.Monad (replicateM)
-
-import Debug.Trace
-
-import GHC.IO.Handle
-import System.Exit
-import System.Process hiding (env)
-
+import Control.Monad (unless)
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Sampler
-import Control.Monad.Bayes.Weighted
-import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Inference.RMSMC
+import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Population
+import Control.Monad.Bayes.Sampler
 import Control.Monad.Bayes.Traced
-
+import Control.Monad.Bayes.Weighted
+import Criterion.Main
+import Criterion.Types
+import GHC.IO.Handle
 -- import NonlinearSSM
 import qualified HMM
-import qualified LogReg
 import qualified LDA
+import qualified LogReg
+import System.Exit
+import System.Process hiding (env)
+import System.Random.MWC (GenIO, createSystemRandom)
 
 -- | Path to the Anglican project with benchmarks.
 anglicanPath :: String
@@ -55,56 +50,53 @@
 
 -- | Insert data into an Anglican model.
 anglicanData :: LeinProc -> Model -> IO ()
-anglicanData lein model = do
-  anglican lein ["(use 'nstools.ns)\n"]
-  anglican lein ["(ns+ " ++ anglicanModelName model ++ ")\n"]
+anglicanData leinProc model = do
+  anglican leinProc ["(use 'nstools.ns)\n"]
+  anglican leinProc ["(ns+ " ++ anglicanModelName model ++ ")\n"]
   case model of
     LR dataset -> do
       let (xs, labels) = unzip dataset
-      anglican lein ["(ns-unmap *ns* 'xs)\n"]
-      anglican lein ["(def xs " ++ clojureShowVector xs ++ ")\n", "nil\n"]
-      anglican lein ["(ns-unmap *ns* 'labels)\n"]
-      anglican lein ["(def labels " ++ clojureBoolVector labels ++ ")\n", "nil\n"]
+      anglican leinProc ["(ns-unmap *ns* 'xs)\n"]
+      anglican leinProc ["(def xs " ++ clojureShowVector xs ++ ")\n", "nil\n"]
+      anglican leinProc ["(ns-unmap *ns* 'labels)\n"]
+      anglican leinProc ["(def labels " ++ clojureBoolVector labels ++ ")\n", "nil\n"]
     HMM observations -> do
-      anglican lein ["(ns-unmap *ns* 'observations)\n"]
-      anglican lein ["(def observations " ++ clojureShowVector observations ++ ")\n", "nil\n"]
+      anglican leinProc ["(ns-unmap *ns* 'observations)\n"]
+      anglican leinProc ["(def observations " ++ clojureShowVector observations ++ ")\n", "nil\n"]
     LDA docs -> do
-      anglican lein ["(ns-unmap *ns* 'docs)\n"]
-      anglican lein ["(def docs " ++ clojureVector (map clojureShowVector docs) ++ ")\n", "nil\n"]
-  anglican lein ["(use 'nstools.ns)\n"]
-  anglican lein ["(ns+ anglican.core)\n"]
+      anglican leinProc ["(ns-unmap *ns* 'docs)\n"]
+      anglican leinProc ["(def docs " ++ clojureVector (map clojureShowVector docs) ++ ")\n", "nil\n"]
+  anglican leinProc ["(use 'nstools.ns)\n"]
+  anglican leinProc ["(ns+ anglican.core)\n"]
 
 anglican :: LeinProc -> [String] -> IO ()
 anglican (LeinProc input output _) cmds = do
   -- execute command
-  mapM (hPutStr input) cmds
+  mapM_ (hPutStr input) cmds
   hFlush input
   -- wait until all samples are produced
   waitForAnglican output
 
 -- | Wait until an Anglican program finishes
 waitForAnglican :: Handle -> IO ()
-waitForAnglican handle = run where
-  run = do
-    l <- hGetLine handle
-    -- traceIO l
-    -- We recognize that Anglican is finished when the repl emits a line "nil"
-    if l == "nil" then
-      return ()
-    else
-      run
+waitForAnglican handle = run
+  where
+    run = do
+      l <- hGetLine handle
+      -- traceIO l
+      -- We recognize that Anglican is finished when the repl emits a line "nil"
+      unless (l == "nil") run
 
 -- | Start a Leiningen process in a directory that contains benchmarks.
 startLein :: IO LeinProc
 startLein = do
-  let setup = (shell "lein repl"){cwd = Just anglicanPath, std_in = CreatePipe, std_out = CreatePipe}
+  let setup = (shell "lein repl") {cwd = Just anglicanPath, std_in = CreatePipe, std_out = CreatePipe}
   (Just input, Just output, _, process) <- createProcess setup
   -- wait until Leiningen starts producing output to make sure it's ready
   _ <- hWaitForInput output (-1) -- wait for output indefinitely
-  let lein = LeinProc input output process
-  anglican lein ["(use 'nstools.ns)\n"]
-  return lein
-
+  let leinProc = LeinProc input output process
+  anglican leinProc ["(use 'nstools.ns)\n"]
+  return leinProc
 
 -- | Path to the WebPPL project with benchmarks.
 webpplPath :: String
@@ -123,27 +115,32 @@
 data Env = Env {rng :: GenIO, lein :: LeinProc}
 
 data ProbProgSys = MonadBayes | Anglican | WebPPL
-  deriving(Show)
+  deriving (Show)
 
-data Model = LR [(Double,Bool)] | HMM [Double] | LDA [[String]]
+data Model = LR [(Double, Bool)] | HMM [Double] | LDA [[String]]
+
 instance Show Model where
   show (LR xs) = "LR" ++ show (length xs)
   show (HMM xs) = "HMM" ++ show (length xs)
   show (LDA xs) = "LDA" ++ show (length $ head xs)
+
 buildModel :: MonadInfer m => Model -> m String
 buildModel (LR dataset) = show <$> LogReg.logisticRegression dataset
 buildModel (HMM dataset) = show <$> HMM.hmm dataset
 buildModel (LDA dataset) = show <$> LDA.lda dataset
+
 modelLength :: Model -> Int
 modelLength (LR xs) = length xs
 modelLength (HMM xs) = length xs
 modelLength (LDA xs) = sum (map length xs)
 
 data Alg = MH Int | SMC Int | RMSMC Int Int
+
 instance Show Alg where
   show (MH n) = "MH" ++ show n
   show (SMC n) = "SMC" ++ show n
   show (RMSMC n t) = "RMSMC" ++ show n ++ "-" ++ show t
+
 runAlg :: Model -> Alg -> SamplerIO String
 runAlg model (MH n) = show <$> prior (mh n (buildModel model))
 runAlg model (SMC n) = show <$> runPopulation (smcSystematic (modelLength model) n (buildModel model))
@@ -151,103 +148,110 @@
 
 prepareBenchmarkable :: GenIO -> ProbProgSys -> Model -> Alg -> Benchmarkable
 prepareBenchmarkable g MonadBayes model alg = nfIO $ sampleIOwith (runAlg model alg) g
-prepareBenchmarkable _ Anglican _ _  = error "Anglican benchmarks not available"
+prepareBenchmarkable _ Anglican _ _ = error "Anglican benchmarks not available"
 prepareBenchmarkable _ WebPPL _ _ = error "WebPPL benchmarks not available"
 
 prepareBenchmark :: Env -> ProbProgSys -> Model -> Alg -> Benchmark
 prepareBenchmark e MonadBayes model alg =
   bench (show MonadBayes ++ sep ++ show model ++ sep ++ show alg) $
-  prepareBenchmarkable (rng e) MonadBayes model alg where
+    prepareBenchmarkable (rng e) MonadBayes model alg
+  where
     sep = "_"
-prepareBenchmark e Anglican model alg = env prepareData (const $ bench name $ whnfIO collect) where
-  name = show Anglican ++ sep ++ show model ++ sep ++ show alg
-  sep = "_"
-  algString (MH n) = "-a lmh -n " ++ show n
-  algString (SMC n) = "-a smc -n " ++ show n
-  algString (RMSMC _ _) = error "Anglican does not support resample-move SMC"
-  prepareData = anglicanData (lein e) model
-  collect = do
-    anglican (lein e) $ ["(time (m! " ++ anglicanModelName model ++ " " ++ algString alg ++ "))\n"]
-prepareBenchmark _ WebPPL model alg = bench name $ whnfIO run where
-  name = show WebPPL ++ sep ++ show model ++ sep ++ show alg
-  sep = "_"
-  algString (MH n) = "--alg MCMC --samples " ++ show n ++ " --rejuv 0 "
-  algString (SMC n) = "--alg SMC --samples " ++ show n ++ " --rejuv 0 "
-  algString (RMSMC n t) = "--alg SMC --samples " ++ show n ++ " --rejuv " ++ show t ++ " "
-  dataString (LR dataset) = let (xs, labels) = unzip dataset in "--xs='" ++ javascriptList xs ++ "' --labels='" ++ javascriptList (map (\b -> if b then 1 else 0) labels) ++ "'"
-  dataString (HMM obs) = "--obs='" ++ javascriptList obs ++ "'"
-  dataString (LDA docs) = unwords $ map (\(i,doc) -> "--doc" ++ show i ++ "='" ++ unwords doc ++ "'") (zip [1..5] docs)
-  run = do
-    let command = "node " ++ webpplModelName model ++ ".js " ++ algString alg ++ dataString model
-    (_, _, _, process) <- createProcess $ (shell command){cwd = Just webpplPath, std_out = NoStream, std_err = NoStream}
-    exitCode <- waitForProcess process
-    case exitCode of
-      ExitSuccess -> return ()
-      ExitFailure i -> error $ "WebPPL terminated with exit code " ++ show i
-
+prepareBenchmark e Anglican model alg = env prepareData (const $ bench name $ whnfIO collect)
+  where
+    name = show Anglican ++ sep ++ show model ++ sep ++ show alg
+    sep = "_"
+    algString (MH n) = "-a lmh -n " ++ show n
+    algString (SMC n) = "-a smc -n " ++ show n
+    algString (RMSMC _ _) = error "Anglican does not support resample-move SMC"
+    prepareData = anglicanData (lein e) model
+    collect =
+      anglican (lein e) ["(time (m! " ++ anglicanModelName model ++ " " ++ algString alg ++ "))\n"]
+prepareBenchmark _ WebPPL model alg = bench name $ whnfIO run
+  where
+    name = show WebPPL ++ sep ++ show model ++ sep ++ show alg
+    sep = "_"
+    algString (MH n) = "--alg MCMC --samples " ++ show n ++ " --rejuv 0 "
+    algString (SMC n) = "--alg SMC --samples " ++ show n ++ " --rejuv 0 "
+    algString (RMSMC n t) = "--alg SMC --samples " ++ show n ++ " --rejuv " ++ show t ++ " "
+    dataString (LR dataset) = let (xs, labels) = unzip dataset in "--xs='" ++ javascriptList xs ++ "' --labels='" ++ javascriptList (map (\b -> if b then (1 :: Int) else 0) labels) ++ "'"
+    dataString (HMM obs) = "--obs='" ++ javascriptList obs ++ "'"
+    dataString (LDA docs) = unwords $ zipWith (\i doc -> "--doc" ++ show (i :: Int) ++ "='" ++ unwords doc ++ "'") [1 .. 5] docs
+    run = do
+      let command = "node " ++ webpplModelName model ++ ".js " ++ algString alg ++ dataString model
+      (_, _, _, process) <- createProcess $ (shell command) {cwd = Just webpplPath, std_out = NoStream, std_err = NoStream}
+      exitCode <- waitForProcess process
+      case exitCode of
+        ExitSuccess -> return ()
+        ExitFailure i -> error $ "WebPPL terminated with exit code " ++ show i
 
 -- | Checks if the requested benchmark is implemented.
 supported :: (ProbProgSys, Model, Alg) -> Bool
 supported (Anglican, _, RMSMC _ _) = False
 supported _ = True
 
-systems = [
-            MonadBayes
-            -- Anglican,
-            -- WebPPL
-          ]
+systems :: [ProbProgSys]
+systems =
+  [ MonadBayes
+    -- Anglican,
+    -- WebPPL
+  ]
 
-lengthBenchmarks e lrData hmmData ldaData = benchmarks where
-  lrLengths = [10] ++ map (*100) [1..10]
-  hmmLengths = [10] ++ map (*100) [1..10]
-  ldaLengths = [5] ++ map (*50) [1..10]
-  models =
-    map (LR . (`take` lrData)) lrLengths ++
-    map (HMM . (`take` hmmData)) hmmLengths ++
-    map (\n -> LDA $ map (take n) ldaData) ldaLengths
-  algs = [
-    MH 100,
-    SMC 100,
-    RMSMC 10 1
-    ]
-  benchmarks = map (uncurry3 (prepareBenchmark e)) $ filter supported xs where
-        uncurry3 f (x,y,z) = f x y z
+lengthBenchmarks :: Env -> [(Double, Bool)] -> [Double] -> [[String]] -> [Benchmark]
+lengthBenchmarks e lrData hmmData ldaData = benchmarks
+  where
+    lrLengths = 10 : map (* 100) [1 .. 10]
+    hmmLengths = 10 : map (* 100) [1 .. 10]
+    ldaLengths = 5 : map (* 50) [1 .. 10]
+    models =
+      map (LR . (`take` lrData)) lrLengths
+        ++ map (HMM . (`take` hmmData)) hmmLengths
+        ++ map (\n -> LDA $ map (take n) ldaData) ldaLengths
+    algs =
+      [ MH 100,
+        SMC 100,
+        RMSMC 10 1
+      ]
+    benchmarks = map (uncurry3 (prepareBenchmark e)) $ filter supported xs
+      where
+        uncurry3 f (x, y, z) = f x y z
         xs = do
           m <- models
           s <- systems
           a <- algs
-          return (s,m,a)
+          return (s, m, a)
 
-samplesBenchmarks e lrData hmmData ldaData = benchmarks where
-  lrLengths = [50]
-  hmmLengths = [20]
-  ldaLengths = [10]
-  models = map (LR . (`take` lrData)) lrLengths ++
-           map (HMM . (`take` hmmData)) hmmLengths ++
-           map (\n -> LDA $ map (take n) ldaData) ldaLengths
-  algs =  map (\x -> MH (100*x)) [1..10] ++ map (\x -> SMC (100*x)) [1..10]
-          ++ map (\x -> RMSMC 10 (10*x)) [1..10]
-  benchmarks = map (uncurry3 (prepareBenchmark e)) $ filter supported xs where
-        uncurry3 f (x,y,z) = f x y z
+samplesBenchmarks :: Env -> [(Double, Bool)] -> [Double] -> [[String]] -> [Benchmark]
+samplesBenchmarks e lrData hmmData ldaData = benchmarks
+  where
+    lrLengths = [50]
+    hmmLengths = [20]
+    ldaLengths = [10]
+    models =
+      map (LR . (`take` lrData)) lrLengths
+        ++ map (HMM . (`take` hmmData)) hmmLengths
+        ++ map (\n -> LDA $ map (take n) ldaData) ldaLengths
+    algs =
+      map (\x -> MH (100 * x)) [1 .. 10] ++ map (\x -> SMC (100 * x)) [1 .. 10]
+        ++ map (\x -> RMSMC 10 (10 * x)) [1 .. 10]
+    benchmarks = map (uncurry3 (prepareBenchmark e)) $ filter supported xs
+      where
+        uncurry3 f (x, y, z) = f x y z
         xs = do
           a <- algs
           s <- systems
           m <- models
-          return (s,m,a)
+          return (s, m, a)
 
 main :: IO ()
 main = do
-
   g <- createSystemRandom
   l <- startLein
   let e = Env g l
-
   lrData <- sampleIOwith (LogReg.syntheticData 1000) g
   hmmData <- sampleIOwith (HMM.syntheticData 1000) g
   ldaData <- sampleIOwith (LDA.syntheticData 5 1000) g
-
-  let configLength = defaultConfig{csvFile = Just "speed-length.csv", rawDataFile = Just "raw.dat"}
+  let configLength = defaultConfig {csvFile = Just "speed-length.csv", rawDataFile = Just "raw.dat"}
   defaultMainWith configLength (lengthBenchmarks e lrData hmmData ldaData)
-
-  let configSamples = defaultConfig{csvFile = Just "speed-samples.csv", rawDataFile = Just "raw.dat"}
+  let configSamples = defaultConfig {csvFile = Just "speed-samples.csv", rawDataFile = Just "raw.dat"}
   defaultMainWith configSamples (samplesBenchmarks e lrData hmmData ldaData)
diff --git a/models/Dice.hs b/models/Dice.hs
--- a/models/Dice.hs
+++ b/models/Dice.hs
@@ -1,5 +1,3 @@
-
-
 module Dice where
 
 -- A toy model for dice rolling from http://dl.acm.org/citation.cfm?id=2804317
@@ -10,23 +8,23 @@
 
 -- | A toss of a six-sided die.
 die :: MonadSample m => m Int
-die = uniformD [1..6]
+die = uniformD [1 .. 6]
 
 -- | A sum of outcomes of n independent tosses of six-sided dice.
 dice :: MonadSample m => Int -> m Int
 dice 1 = die
-dice n = liftM2 (+) die (dice (n-1))
+dice n = liftM2 (+) die (dice (n -1))
 
 -- | Toss of two dice where the output is greater than 4.
-dice_hard :: MonadInfer m => m Int
-dice_hard = do
+diceHard :: MonadInfer m => m Int
+diceHard = do
   result <- dice 2
   condition (result > 4)
   return result
 
 -- | Toss of two dice with an artificial soft constraint.
-dice_soft :: MonadInfer m => m Int
-dice_soft = do
+diceSoft :: MonadInfer m => m Int
+diceSoft = do
   result <- dice 2
   score (1 / fromIntegral result)
   return result
diff --git a/models/HMM.hs b/models/HMM.hs
--- a/models/HMM.hs
+++ b/models/HMM.hs
@@ -1,55 +1,72 @@
--- HMM from Anglican (https://bitbucket.org/probprog/anglican-white-paper)
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TypeFamilies #-}
 
-{-# LANGUAGE
- FlexibleContexts,
- TypeFamilies
- #-}
+-- HMM from Anglican (https://bitbucket.org/probprog/anglican-white-paper)
 
-module HMM (
-  values,
-  hmm,
-  syntheticData
-  ) where
+module HMM
+  ( values,
+    hmm,
+    syntheticData,
+  )
+where
 
 --Hidden Markov Models
 
 import Control.Monad (replicateM)
-import Data.Vector (fromList)
-
 import Control.Monad.Bayes.Class
+import Data.Vector (fromList)
 
 -- | Observed values
 values :: [Double]
-values = [0.9,0.8,0.7,0,-0.025,-5,-2,-0.1,0,
-          0.13,0.45,6,0.2,0.3,-1,-1]
+values =
+  [ 0.9,
+    0.8,
+    0.7,
+    0,
+    -0.025,
+    -5,
+    -2,
+    -0.1,
+    0,
+    0.13,
+    0.45,
+    6,
+    0.2,
+    0.3,
+    -1,
+    -1
+  ]
 
 -- | The transition model.
 trans :: MonadSample m => Int -> m Int
 trans 0 = categorical $ fromList [0.1, 0.4, 0.5]
 trans 1 = categorical $ fromList [0.2, 0.6, 0.2]
-trans 2 = categorical $ fromList [0.15,0.7,0.15]
+trans 2 = categorical $ fromList [0.15, 0.7, 0.15]
+trans _ = error "unreachable"
 
 -- | The emission model.
 emissionMean :: Int -> Double
-emissionMean x = mean x where
-  mean 0 = -1
-  mean 1 = 1
-  mean 2 = 0
+emissionMean 0 = -1
+emissionMean 1 = 1
+emissionMean 2 = 0
+emissionMean _ = error "unreachable"
 
 -- | Initial state distribution
 start :: MonadSample m => m Int
-start = uniformD [0,1,2]
+start = uniformD [0, 1, 2]
 
 -- | Example HMM from http://dl.acm.org/citation.cfm?id=2804317
 hmm :: (MonadInfer m) => [Double] -> m [Int]
-hmm dataset = f dataset (const . return) where
-  expand x y = do
-    x' <- trans x
-    factor $ normalPdf (emissionMean x') 1 y
-    return x'
-  f [] k = start >>= k []
-  f (y:ys) k = f ys (\xs x -> expand x y >>= k (x:xs))
+hmm dataset = f dataset (const . return)
+  where
+    expand x y = do
+      x' <- trans x
+      factor $ normalPdf (emissionMean x') 1 y
+      return x'
+    f [] k = start >>= k []
+    f (y : ys) k = f ys (\xs x -> expand x y >>= k (x : xs))
 
 syntheticData :: MonadSample m => Int -> m [Double]
-syntheticData n = replicateM n syntheticPoint where
-    syntheticPoint = uniformD [0,1,2]
+syntheticData n = replicateM n syntheticPoint
+  where
+    syntheticPoint = uniformD [0, 1, 2]
diff --git a/models/LDA.hs b/models/LDA.hs
--- a/models/LDA.hs
+++ b/models/LDA.hs
@@ -3,54 +3,52 @@
 
 module LDA where
 
-import Numeric.Log
 import qualified Control.Monad as List (replicateM)
-import Data.Vector as V hiding (length, zip, mapM, mapM_)
+import Control.Monad.Bayes.Class
 import qualified Data.Map as Map
+import Data.Vector as V hiding (length, mapM, mapM_, zip)
+import Numeric.Log
 
-import Control.Monad.Bayes.Class
+vocabulary :: [String]
+vocabulary = ["bear", "wolf", "python", "prolog"]
 
-vocabluary :: [String]
-vocabluary = ["bear", "wolf", "python", "prolog"]
 topics :: [String]
 topics = ["topic1", "topic2"]
 
-docs :: [[String]]
-docs = [
-  words "bear wolf bear wolf bear wolf python wolf bear wolf",
-  words "python prolog python prolog python prolog python prolog python prolog",
-  words "bear wolf bear wolf bear wolf bear wolf bear wolf",
-  words "python prolog python prolog python prolog python prolog python prolog",
-  words "bear wolf bear python bear wolf bear wolf bear wolf"
+documents :: [[String]]
+documents =
+  [ words "bear wolf bear wolf bear wolf python wolf bear wolf",
+    words "python prolog python prolog python prolog python prolog python prolog",
+    words "bear wolf bear wolf bear wolf bear wolf bear wolf",
+    words "python prolog python prolog python prolog python prolog python prolog",
+    words "bear wolf bear python bear wolf bear wolf bear wolf"
   ]
 
-word_dist_prior :: MonadSample m => m (Vector Double)
-word_dist_prior = dirichlet $ V.replicate (length vocabluary) 1
+wordDistPrior :: MonadSample m => m (Vector Double)
+wordDistPrior = dirichlet $ V.replicate (length vocabulary) 1
 
-topic_dist_prior :: MonadSample m => m (Vector Double)
-topic_dist_prior = dirichlet $ V.replicate (length topics) 1
+topicDistPrior :: MonadSample m => m (Vector Double)
+topicDistPrior = dirichlet $ V.replicate (length topics) 1
 
-word_index :: Map.Map String Int
-word_index = Map.fromList $ zip vocabluary [0..]
+wordIndex :: Map.Map String Int
+wordIndex = Map.fromList $ zip vocabulary [0 ..]
 
 lda :: MonadInfer m => [[String]] -> m [Int]
 lda docs = do
   word_dist_for_topic <- do
-    ts <- mapM (const word_dist_prior) [0 .. length topics]
+    ts <- mapM (const wordDistPrior) [0 .. length topics]
     return $ Map.fromList $ zip [0 .. length topics] ts
-
   let obs doc = do
-        topic_dist <- fmap categorical topic_dist_prior
+        topic_dist <- fmap categorical topicDistPrior
         let f word = do
               topic <- topic_dist
-              factor $ (Exp . log) $ (word_dist_for_topic Map.! topic) V.! (word_index Map.! word)
+              factor $ (Exp . log) $ (word_dist_for_topic Map.! topic) V.! (wordIndex Map.! word)
         mapM_ f doc
-
   mapM_ obs docs
-
   -- return samples since Discrete is not NFData
   mapM (categorical . snd) $ Map.toList word_dist_for_topic
 
 syntheticData :: MonadSample m => Int -> Int -> m [[String]]
-syntheticData d w = List.replicateM d (List.replicateM w syntheticWord) where
-  syntheticWord = uniformD vocabluary
+syntheticData d w = List.replicateM d (List.replicateM w syntheticWord)
+  where
+    syntheticWord = uniformD vocabulary
diff --git a/models/LogReg.hs b/models/LogReg.hs
--- a/models/LogReg.hs
+++ b/models/LogReg.hs
@@ -4,9 +4,8 @@
 module LogReg where
 
 import Control.Monad (replicateM)
-
-import Numeric.Log
 import Control.Monad.Bayes.Class
+import Numeric.Log
 
 xs :: [Double]
 xs = [-10, -5, 2, 6, 10]
@@ -28,8 +27,9 @@
   sigmoid 8
 
 syntheticData :: MonadSample m => Int -> m [(Double, Bool)]
-syntheticData n = replicateM n syntheticPoint where
+syntheticData n = replicateM n syntheticPoint
+  where
     syntheticPoint = do
       x <- uniform (-1) 1
       label <- bernoulli 0.5
-      return (x,label)
+      return (x, label)
diff --git a/models/NonlinearSSM.hs b/models/NonlinearSSM.hs
--- a/models/NonlinearSSM.hs
+++ b/models/NonlinearSSM.hs
@@ -12,42 +12,46 @@
   let sigmaY = 1 / sqrt precY
   return (sigmaX, sigmaY)
 
+mean :: Double -> Int -> Double
+mean x n = 0.5 * x + 25 * x / (1 + x * x) + 8 * cos (1.2 * fromIntegral n)
+
 -- | A nonlinear series model from Doucet et al. (2000)
 -- "On sequential Monte Carlo sampling methods" section VI.B
-model :: (MonadInfer m)
-      => [Double]  -- ^ observed data
-      -> (Double, Double) -- ^ prior on the parameters
-      -> m [Double] -- ^ list of latent states from t=1
+model ::
+  (MonadInfer m) =>
+  -- | observed data
+  [Double] ->
+  -- | prior on the parameters
+  (Double, Double) ->
+  -- | list of latent states from t=1
+  m [Double]
 model obs (sigmaX, sigmaY) = do
   let sq x = x * x
       simulate [] _ acc = return acc
-      simulate (y:ys) x acc = do
+      simulate (y : ys) x acc = do
         let n = length acc
-        let mean = 0.5 * x + 25 * x / (1 + sq x) +
-                   8 * cos (1.2 * fromIntegral n)
-        x' <- normal mean sigmaX
+        x' <- normal (mean x n) sigmaX
         factor $ normalPdf (sq x' / 20) sigmaY y
-        simulate ys x' (x':acc)
-
+        simulate ys x' (x' : acc)
   x0 <- normal 0 (sqrt 5)
   xs <- simulate obs x0 []
   return $ reverse xs
 
-generateData :: MonadSample m
-      => Int  -- ^ T
-      -> m [(Double,Double)] -- ^ list of latent and observable states from t=1
+generateData ::
+  MonadSample m =>
+  -- | T
+  Int ->
+  -- | list of latent and observable states from t=1
+  m [(Double, Double)]
 generateData t = do
   (sigmaX, sigmaY) <- param
   let sq x = x * x
       simulate 0 _ acc = return acc
       simulate k x acc = do
         let n = length acc
-        let mean = 0.5 * x + 25 * x / (1 + sq x) +
-                   8 * cos (1.2 * fromIntegral n)
-        x' <- normal mean sigmaX
+        x' <- normal (mean x n) sigmaX
         y' <- normal (sq x' / 20) sigmaY
-        simulate (k-1) x' ((x',y'):acc)
-
+        simulate (k -1) x' ((x', y') : acc)
   x0 <- normal 0 (sqrt 5)
   xys <- simulate t x0 []
   return $ reverse xys
diff --git a/models/Sprinkler.hs b/models/Sprinkler.hs
--- a/models/Sprinkler.hs
+++ b/models/Sprinkler.hs
@@ -1,17 +1,17 @@
 module Sprinkler where
 
 import Control.Monad (when)
-
 import Control.Monad.Bayes.Class
 
 hard :: MonadInfer m => m Bool
 hard = do
   rain <- bernoulli 0.3
   sprinkler <- bernoulli $ if rain then 0.1 else 0.4
-  wet <- bernoulli $ case (rain,sprinkler) of (True,True) -> 0.98
-                                              (True,False) -> 0.8
-                                              (False,True) -> 0.9
-                                              (False,False) -> 0.0
+  wet <- bernoulli $ case (rain, sprinkler) of
+    (True, True) -> 0.98
+    (True, False) -> 0.8
+    (False, True) -> 0.9
+    (False, False) -> 0.0
   condition (not wet)
   return rain
 
diff --git a/monad-bayes.cabal b/monad-bayes.cabal
--- a/monad-bayes.cabal
+++ b/monad-bayes.cabal
@@ -1,141 +1,183 @@
-cabal-version:       2.0
-name:                monad-bayes
-version:             0.1.0.0
-synopsis:            A library for probabilistic programming.
-description:         A library for probabilistic programming using probability monads. The emphasis is on composition of inference algorithms implemented in terms of monad transformers.
-homepage:            http://github.com/tweag/monad-bayes#readme
-license:             MIT
-license-file:        LICENSE
-author:              Adam Scibior <adscib@gmail.com>
-maintainer:          leonhard.markert@tweag.io
-copyright:           2015-2020 Adam Scibior
-bug-reports:         https://github.com/tweag/monad-bayes/issues
-stability:           experimental
-category:            Statistics
-build-type:          Simple
+cabal-version:      2.0
+name:               monad-bayes
+version:            0.1.1.0
+license:            MIT
+license-file:       LICENSE.md
+copyright:          2015-2020 Adam Scibior
+maintainer:         leonhard.markert@tweag.io
+author:             Adam Scibior <adscib@gmail.com>
+stability:          experimental
+tested-with:        ghc ==8.4.4 ghc ==8.6.5 ghc ==8.8.1
+homepage:           http://github.com/tweag/monad-bayes#readme
+bug-reports:        https://github.com/tweag/monad-bayes/issues
+synopsis:           A library for probabilistic programming.
+description:
+    A library for probabilistic programming using probability monads. The
+    emphasis is on composition of inference algorithms implemented in
+    terms of monad transformers.
 
-extra-source-files:
-  CHANGELOG.md
+category:           Statistics
+build-type:         Simple
+extra-source-files: CHANGELOG.md
 
-executable example
-  hs-source-dirs:      benchmark, models
-  main-is:             Single.hs
-  build-depends:       base
-                     , monad-bayes
-                     , log-domain
-                     , vector
-                     , containers
-                     , mwc-random
-                     , time
-                     , optparse-applicative
-  default-language:    Haskell2010
-  other-modules:       LogReg
-                     , HMM
-                     , LDA
-                     , Dice
+source-repository head
+    type:     git
+    location: https://github.com/tweag/monad-bayes.git
 
+flag dev
+    description: Turn on development settings.
+    default:     False
+    manual:      True
+
 library
-  hs-source-dirs:      src
-  exposed-modules:     Control.Monad.Bayes.Class
-                     , Control.Monad.Bayes.Free
-                     , Control.Monad.Bayes.Sampler
-                     , Control.Monad.Bayes.Weighted
-                     , Control.Monad.Bayes.Sequential
-                     , Control.Monad.Bayes.Population
-                     , Control.Monad.Bayes.Traced.Static
-                     , Control.Monad.Bayes.Traced.Dynamic
-                     , Control.Monad.Bayes.Traced.Basic
-                     , Control.Monad.Bayes.Traced
-                     , Control.Monad.Bayes.Enumerator
-                     , Control.Monad.Bayes.Helpers
-                     , Control.Monad.Bayes.Inference.SMC
-                     , Control.Monad.Bayes.Inference.RMSMC
-                     , Control.Monad.Bayes.Inference.PMMH
-                     , Control.Monad.Bayes.Inference.SMC2
-  other-modules:       Control.Monad.Bayes.Traced.Common
-  -- See MAINTAINERS.md for when and how to update the version bounds.
-  build-depends:       base >= 4.10.1 && < 4.14
-                     , containers >= 0.5.10 && < 0.7
-                     , free >= 5.0.2 && < 5.2
-                     , ieee754 >= 0.8.0 && < 0.9
-                     , log-domain >= 0.12 && < 0.14
-                     , math-functions >= 0.2.1 && < 0.4
-                     , monad-coroutine >= 0.9.0 && < 0.10
-                     , mtl >= 2.2.2 && < 2.3
-                     , mwc-random >= 0.13.6 && < 0.15
-                     , safe >= 0.3.17 && < 0.4
-                     , statistics >= 0.14.0 && < 0.16
-                     , transformers >= 0.5.2 && < 0.6
-                     , vector >= 0.12.0 && < 0.13
-  ghc-options:         -Wall -fno-warn-redundant-constraints
-  default-language:    Haskell2010
-  default-extensions:  RankNTypes
-                     , GeneralizedNewtypeDeriving
-                     , StandaloneDeriving
-                     , TypeFamilies
-                     , FlexibleContexts
-                     , FlexibleInstances
-                     , TupleSections
-                     , MultiParamTypeClasses
-                     , GADTs
-  other-extensions:    ScopedTypeVariables
-                     , DeriveFunctor
+    exposed-modules:
+        Control.Monad.Bayes.Class
+        Control.Monad.Bayes.Enumerator
+        Control.Monad.Bayes.Free
+        Control.Monad.Bayes.Helpers
+        Control.Monad.Bayes.Inference.PMMH
+        Control.Monad.Bayes.Inference.RMSMC
+        Control.Monad.Bayes.Inference.SMC
+        Control.Monad.Bayes.Inference.SMC2
+        Control.Monad.Bayes.Population
+        Control.Monad.Bayes.Sampler
+        Control.Monad.Bayes.Sequential
+        Control.Monad.Bayes.Traced
+        Control.Monad.Bayes.Traced.Basic
+        Control.Monad.Bayes.Traced.Dynamic
+        Control.Monad.Bayes.Traced.Static
+        Control.Monad.Bayes.Weighted
 
+    hs-source-dirs:     src
+    other-modules:      Control.Monad.Bayes.Traced.Common
+    default-language:   Haskell2010
+    default-extensions:
+        MultiParamTypeClasses RankNTypes FlexibleContexts FlexibleInstances
+        GeneralizedNewtypeDeriving TypeFamilies StandaloneDeriving GADTs
+        TupleSections
+
+    other-extensions:   ScopedTypeVariables DeriveFunctor
+    build-depends:
+        base >=4.11 && <4.14,
+        containers >=0.5.10 && <0.7,
+        free >=5.0.2 && <5.2,
+        ieee754 >=0.8.0 && <0.9,
+        log-domain >=0.12 && <0.14,
+        math-functions >=0.2.1 && <0.4,
+        monad-coroutine >=0.9.0 && <0.10,
+        mtl >=2.2.2 && <2.3,
+        mwc-random >=0.13.6 && <0.15,
+        safe >=0.3.17 && <0.4,
+        statistics >=0.14.0 && <0.16,
+        transformers >=0.5.2 && <0.6,
+        vector >=0.12.0 && <0.13
+
+    if flag(dev)
+        ghc-options:
+            -Wall -Wcompat -Wincomplete-record-updates
+            -Wincomplete-uni-patterns -Wnoncanonical-monad-instances
+
+    else
+        ghc-options: -Wall
+
+executable example
+    main-is:          Single.hs
+    hs-source-dirs:   benchmark models
+    other-modules:
+        Dice
+        HMM
+        LDA
+        LogReg
+
+    default-language: Haskell2010
+    build-depends:
+        base -any,
+        containers -any,
+        log-domain -any,
+        monad-bayes -any,
+        mwc-random -any,
+        optparse-applicative -any,
+        time -any,
+        vector -any
+
+    if flag(dev)
+        ghc-options:
+            -Wall -Wcompat -Wincomplete-record-updates
+            -Wincomplete-uni-patterns -Wnoncanonical-monad-instances
+
+    else
+        ghc-options: -Wall
+
 test-suite monad-bayes-test
-  type:                exitcode-stdio-1.0
-  hs-source-dirs:      test, models
-  main-is:             Spec.hs
-  build-depends:       base
-                     , monad-bayes
-                     , hspec
-                     , QuickCheck
-                     , ieee754
-                     , mtl
-                     , math-functions
-                     , transformers
-                     , log-domain
-                     , vector
-  ghc-options:         -Wall -threaded -rtsopts "-with-rtsopts=-N -M1g -K1g"
-  default-language:    Haskell2010
-  other-modules:       Sprinkler,
-                       TestEnumerator,
-                       TestPopulation,
-                       TestInference,
-                       TestSequential,
-                       TestWeighted
+    type:             exitcode-stdio-1.0
+    main-is:          Spec.hs
+    hs-source-dirs:   test models
+    other-modules:
+        Sprinkler
+        TestEnumerator
+        TestInference
+        TestPopulation
+        TestSequential
+        TestWeighted
 
+    default-language: Haskell2010
+    build-depends:
+        base -any,
+        QuickCheck -any,
+        hspec -any,
+        ieee754 -any,
+        log-domain -any,
+        math-functions -any,
+        monad-bayes -any,
+        mtl -any,
+        transformers -any,
+        vector -any
+
+    if flag(dev)
+        ghc-options:
+            -Wall -Wcompat -Wincomplete-record-updates
+            -Wincomplete-uni-patterns -Wnoncanonical-monad-instances
+
+    else
+        ghc-options: -Wall
+
 benchmark ssm-bench
-  type:                exitcode-stdio-1.0
-  hs-source-dirs:      models
-                     , benchmark
-  build-depends:       base
-                     , monad-bayes
-  default-language:    Haskell2010
-  default-extensions:  RankNTypes
-  main-is:             SSM.hs
-  other-modules:       NonlinearSSM
+    type:               exitcode-stdio-1.0
+    main-is:            SSM.hs
+    hs-source-dirs:     models benchmark
+    other-modules:      NonlinearSSM
+    default-language:   Haskell2010
+    default-extensions: RankNTypes
+    build-depends:
+        base -any,
+        monad-bayes -any
 
 benchmark speed-bench
-  type:                exitcode-stdio-1.0
-  hs-source-dirs:      models
-                     , benchmark
-  build-depends:       base
-                     , monad-bayes
-                     , criterion
-                     , abstract-par
-                     , log-domain
-                     , vector
-                     , mwc-random
-                     , containers
-                     , process
-  ghc-options:         -rtsopts "-with-rtsopts=-M1g -K1g"
-  default-language:    Haskell2010
-  default-extensions:  RankNTypes
-  main-is:             Speed.hs
-  other-modules:       LogReg
-                     , HMM
-                     , LDA
+    type:               exitcode-stdio-1.0
+    main-is:            Speed.hs
+    hs-source-dirs:     models benchmark
+    other-modules:
+        HMM
+        LDA
+        LogReg
 
-source-repository head
-  type:     git
-  location: https://github.com/tweag/monad-bayes.git
+    default-language:   Haskell2010
+    default-extensions: RankNTypes
+    build-depends:
+        base -any,
+        abstract-par -any,
+        containers -any,
+        criterion -any,
+        log-domain -any,
+        monad-bayes -any,
+        mwc-random -any,
+        process -any,
+        vector -any
+
+    if flag(dev)
+        ghc-options:
+            -Wall -Wcompat -Wincomplete-record-updates
+            -Wincomplete-uni-patterns -Wnoncanonical-monad-instances
+
+    else
+        ghc-options: -Wall
diff --git a/src/Control/Monad/Bayes/Class.hs b/src/Control/Monad/Bayes/Class.hs
--- a/src/Control/Monad/Bayes/Class.hs
+++ b/src/Control/Monad/Bayes/Class.hs
@@ -1,141 +1,160 @@
-{-|
-Module      : Control.Monad.Bayes.Class
-Description : Types for probabilistic modelling
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-This module defines 'MonadInfer', which can be used to represent a simple model
-like the following:
+-- |
+-- Module      : Control.Monad.Bayes.Class
+-- Description : Types for probabilistic modelling
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- This module defines 'MonadInfer', which can be used to represent a simple model
+-- like the following:
+--
+-- @
+-- import Control.Monad (when)
+-- import Control.Monad.Bayes.Class
+--
+-- model :: MonadInfer m => m Bool
+-- model = do
+--   rain <- bernoulli 0.3
+--   sprinkler <-
+--     bernoulli $
+--     if rain
+--       then 0.1
+--       else 0.4
+--   let wetProb =
+--     case (rain, sprinkler) of
+--       (True,  True)  -> 0.98
+--       (True,  False) -> 0.80
+--       (False, True)  -> 0.90
+--       (False, False) -> 0.00
+--   score wetProb
+--   return rain
+-- @
+module Control.Monad.Bayes.Class
+  ( MonadSample,
+    random,
+    uniform,
+    normal,
+    gamma,
+    beta,
+    bernoulli,
+    categorical,
+    logCategorical,
+    uniformD,
+    geometric,
+    poisson,
+    dirichlet,
+    MonadCond,
+    score,
+    factor,
+    condition,
+    MonadInfer,
+    discrete,
+    normalPdf,
+  )
+where
 
-@
 import Control.Monad (when)
-import Control.Monad.Bayes.Class
-
-model :: MonadInfer m => m Bool
-model = do
-  rain <- bernoulli 0.3
-  sprinkler <-
-    bernoulli $
-    if rain
-      then 0.1
-      else 0.4
-  let wetProb =
-    case (rain, sprinkler) of
-      (True,  True)  -> 0.98
-      (True,  False) -> 0.80
-      (False, True)  -> 0.90
-      (False, False) -> 0.00
-  score wetProb
-  return rain
-@
--}
-
-module Control.Monad.Bayes.Class (
-  MonadSample,
-  random,
-  uniform,
-  normal,
-  gamma,
-  beta,
-  bernoulli,
-  categorical,
-  logCategorical,
-  uniformD,
-  geometric,
-  poisson,
-  dirichlet,
-  MonadCond,
-  score,
-  factor,
-  condition,
-  MonadInfer,
-  normalPdf
-) where
-
 import Control.Monad.Trans.Class
+import Control.Monad.Trans.Cont
 import Control.Monad.Trans.Identity
+import Control.Monad.Trans.List
 import Control.Monad.Trans.Maybe
+import Control.Monad.Trans.RWS hiding (tell)
+import Control.Monad.Trans.Reader
 import Control.Monad.Trans.State
 import Control.Monad.Trans.Writer
-import Control.Monad.Trans.Reader
-import Control.Monad.Trans.RWS hiding (tell)
-import Control.Monad.Trans.List
-import Control.Monad.Trans.Cont
-
+import qualified Data.Vector as V
+import Data.Vector.Generic as VG
+import Numeric.Log
 import Statistics.Distribution
-import Statistics.Distribution.Uniform (uniformDistr)
-import Statistics.Distribution.Normal (normalDistr)
-import Statistics.Distribution.Gamma (gammaDistr)
 import Statistics.Distribution.Beta (betaDistr)
+import Statistics.Distribution.Gamma (gammaDistr)
 import Statistics.Distribution.Geometric (geometric0)
+import Statistics.Distribution.Normal (normalDistr)
 import qualified Statistics.Distribution.Poisson as Poisson
-
-import Numeric.Log
-
-import Data.Vector.Generic as VG
-import qualified Data.Vector as V
-import Control.Monad (when)
+import Statistics.Distribution.Uniform (uniformDistr)
 
 -- | Monads that can draw random variables.
 class Monad m => MonadSample m where
   -- | Draw from a uniform distribution.
-  random :: m Double -- ^ \(\sim \mathcal{U}(0, 1)\)
+  random ::
+    -- | \(\sim \mathcal{U}(0, 1)\)
+    m Double
 
   -- | Draw from a uniform distribution.
   uniform ::
-       Double -- ^ lower bound a
-    -> Double -- ^ upper bound b
-    -> m Double -- ^ \(\sim \mathcal{U}(a, b)\).
+    -- | lower bound a
+    Double ->
+    -- | upper bound b
+    Double ->
+    -- | \(\sim \mathcal{U}(a, b)\).
+    m Double
   uniform a b = draw (uniformDistr a b)
 
   -- | Draw from a normal distribution.
   normal ::
-       Double -- ^ mean μ
-    -> Double -- ^ standard deviation σ
-    -> m Double -- ^ \(\sim \mathcal{N}(\mu, \sigma^2)\)
+    -- | mean μ
+    Double ->
+    -- | standard deviation σ
+    Double ->
+    -- | \(\sim \mathcal{N}(\mu, \sigma^2)\)
+    m Double
   normal m s = draw (normalDistr m s)
 
   -- | Draw from a gamma distribution.
   gamma ::
-       Double -- ^ shape k
-    -> Double -- ^ scale θ
-    -> m Double -- ^ \(\sim \Gamma(k, \theta)\)
+    -- | shape k
+    Double ->
+    -- | scale θ
+    Double ->
+    -- | \(\sim \Gamma(k, \theta)\)
+    m Double
   gamma shape scale = draw (gammaDistr shape scale)
 
   -- | Draw from a beta distribution.
   beta ::
-       Double -- ^ shape α
-    -> Double -- ^ shape β
-    -> m Double -- ^ \(\sim \mathrm{Beta}(\alpha, \beta)\)
+    -- | shape α
+    Double ->
+    -- | shape β
+    Double ->
+    -- | \(\sim \mathrm{Beta}(\alpha, \beta)\)
+    m Double
   beta a b = draw (betaDistr a b)
 
   -- | Draw from a Bernoulli distribution.
   bernoulli ::
-       Double -- ^ probability p
-    -> m Bool -- ^ \(\sim \mathrm{B}(1, p)\)
+    -- | probability p
+    Double ->
+    -- | \(\sim \mathrm{B}(1, p)\)
+    m Bool
   bernoulli p = fmap (< p) random
 
   -- | Draw from a categorical distribution.
   categorical ::
-       Vector v Double
-    => v Double -- ^ event probabilities
-    -> m Int -- ^ outcome category
+    Vector v Double =>
+    -- | event probabilities
+    v Double ->
+    -- | outcome category
+    m Int
   categorical ps = fromPMF (ps !)
 
   -- | Draw from a categorical distribution in the log domain.
   logCategorical ::
-       (Vector v (Log Double), Vector v Double)
-    => v (Log Double) -- ^ event probabilities
-    -> m Int -- ^ outcome category
+    (Vector v (Log Double), Vector v Double) =>
+    -- | event probabilities
+    v (Log Double) ->
+    -- | outcome category
+    m Int
   logCategorical = categorical . VG.map (exp . ln)
 
   -- | Draw from a discrete uniform distribution.
   uniformD ::
-       [a] -- ^ observable outcomes @xs@
-    -> m a -- ^ \(\sim \mathcal{U}\{\mathrm{xs}\}\)
+    -- | observable outcomes @xs@
+    [a] ->
+    -- | \(\sim \mathcal{U}\{\mathrm{xs}\}\)
+    m a
   uniformD xs = do
     let n = Prelude.length xs
     i <- categorical $ V.replicate n (1 / fromIntegral n)
@@ -143,21 +162,27 @@
 
   -- | Draw from a geometric distribution.
   geometric ::
-       Double -- ^ success rate p
-    -> m Int -- ^ \(\sim\) number of failed Bernoulli trials with success probability p before first success
+    -- | success rate p
+    Double ->
+    -- | \(\sim\) number of failed Bernoulli trials with success probability p before first success
+    m Int
   geometric = discrete . geometric0
 
   -- | Draw from a Poisson distribution.
   poisson ::
-       Double -- ^ parameter λ
-    -> m Int -- ^ \(\sim \mathrm{Pois}(\lambda)\)
+    -- | parameter λ
+    Double ->
+    -- | \(\sim \mathrm{Pois}(\lambda)\)
+    m Int
   poisson = discrete . Poisson.poisson
 
   -- | Draw from a Dirichlet distribution.
   dirichlet ::
-       Vector v Double
-    => v Double -- ^ concentration parameters @as@
-    -> m (v Double) -- ^ \(\sim \mathrm{Dir}(\mathrm{as})\)
+    Vector v Double =>
+    -- | concentration parameters @as@
+    v Double ->
+    -- | \(\sim \mathrm{Dir}(\mathrm{as})\)
+    m (v Double)
   dirichlet as = do
     xs <- VG.mapM (`gamma` 1) as
     let s = VG.sum xs
@@ -172,13 +197,14 @@
 -- | Draw from a discrete distribution using a sequence of draws from
 -- Bernoulli.
 fromPMF :: MonadSample m => (Int -> Double) -> m Int
-fromPMF p = f 0 1 where
-  f i r = do
-    when (r < 0) $ error "fromPMF: total PMF above 1"
-    let q = p i
-    when (q < 0 || q > 1) $ error "fromPMF: invalid probability value"
-    b <- bernoulli (q / r)
-    if b then pure i else f (i+1) (r-q)
+fromPMF p = f 0 1
+  where
+    f i r = do
+      when (r < 0) $ error "fromPMF: total PMF above 1"
+      let q = p i
+      when (q < 0 || q > 1) $ error "fromPMF: invalid probability value"
+      b <- bernoulli (q / r)
+      if b then pure i else f (i + 1) (r - q)
 
 -- | Draw from a discrete distributions using the probability mass function.
 discrete :: (DiscreteDistr d, MonadSample m) => d -> m Int
@@ -188,14 +214,16 @@
 class Monad m => MonadCond m where
   -- | Record a likelihood.
   score ::
-       Log Double -- ^ likelihood of the execution path
-    -> m ()
+    -- | likelihood of the execution path
+    Log Double ->
+    m ()
 
 -- | Synonym for 'score'.
 factor ::
-     MonadCond m
-  => Log Double -- ^ likelihood of the execution path
-  -> m ()
+  MonadCond m =>
+  -- | likelihood of the execution path
+  Log Double ->
+  m ()
 factor = score
 
 -- | Hard conditioning.
@@ -207,10 +235,14 @@
 
 -- | Probability density function of the normal distribution.
 normalPdf ::
-     Double -- ^ mean μ
-  -> Double -- ^ standard deviation σ
-  -> Double -- ^ sample x
-  -> Log Double -- ^ relative likelihood of observing sample x in \(\mathcal{N}(\mu, \sigma^2)\)
+  -- | mean μ
+  Double ->
+  -- | standard deviation σ
+  Double ->
+  -- | sample x
+  Double ->
+  -- | relative likelihood of observing sample x in \(\mathcal{N}(\mu, \sigma^2)\)
+  Log Double
 normalPdf mu sigma x = Exp $ logDensity (normalDistr mu sigma) x
 
 ----------------------------------------------------------------------------
@@ -225,7 +257,6 @@
 
 instance MonadInfer m => MonadInfer (IdentityT m)
 
-
 instance MonadSample m => MonadSample (MaybeT m) where
   random = lift random
 
@@ -234,7 +265,6 @@
 
 instance MonadInfer m => MonadInfer (MaybeT m)
 
-
 instance MonadSample m => MonadSample (ReaderT r m) where
   random = lift random
   bernoulli = lift . bernoulli
@@ -244,7 +274,6 @@
 
 instance MonadInfer m => MonadInfer (ReaderT r m)
 
-
 instance (Monoid w, MonadSample m) => MonadSample (WriterT w m) where
   random = lift random
   bernoulli = lift . bernoulli
@@ -255,7 +284,6 @@
 
 instance (Monoid w, MonadInfer m) => MonadInfer (WriterT w m)
 
-
 instance MonadSample m => MonadSample (StateT s m) where
   random = lift random
   bernoulli = lift . bernoulli
@@ -266,7 +294,6 @@
 
 instance MonadInfer m => MonadInfer (StateT s m)
 
-
 instance (MonadSample m, Monoid w) => MonadSample (RWST r w s m) where
   random = lift random
 
@@ -275,7 +302,6 @@
 
 instance (MonadInfer m, Monoid w) => MonadInfer (RWST r w s m)
 
-
 instance MonadSample m => MonadSample (ListT m) where
   random = lift random
   bernoulli = lift . bernoulli
@@ -285,7 +311,6 @@
   score = lift . score
 
 instance MonadInfer m => MonadInfer (ListT m)
-
 
 instance MonadSample m => MonadSample (ContT r m) where
   random = lift random
diff --git a/src/Control/Monad/Bayes/Enumerator.hs b/src/Control/Monad/Bayes/Enumerator.hs
--- a/src/Control/Monad/Bayes/Enumerator.hs
+++ b/src/Control/Monad/Bayes/Enumerator.hs
@@ -1,16 +1,13 @@
-{-|
-Module      : Control.Monad.Bayes.Enumerator
-Description : Exhaustive enumeration of discrete random variables
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Enumerator (
-    Enumerator,
+-- |
+-- Module      : Control.Monad.Bayes.Enumerator
+-- Description : Exhaustive enumeration of discrete random variables
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Enumerator
+  ( Enumerator,
     logExplicit,
     explicit,
     evidence,
@@ -18,32 +15,31 @@
     compact,
     enumerate,
     expectation,
-    normalForm
-            ) where
+    normalForm,
+  )
+where
 
-import Data.AEq (AEq, (===), (~==))
-import Control.Applicative (Applicative, Alternative)
-import Control.Monad (MonadPlus)
+import Control.Applicative (Alternative)
 import Control.Arrow (second)
+import Control.Monad (MonadPlus)
+import Control.Monad.Bayes.Class
+import Control.Monad.Trans.Writer
+import Data.AEq ((===), AEq, (~==))
 import qualified Data.Map as Map
+import Data.Maybe
+import Data.Monoid
 import qualified Data.Vector.Generic as V
 import Numeric.Log as Log
-import Control.Monad.Trans.Writer
-import Data.Monoid
-import Data.Maybe
 
-import Control.Monad.Bayes.Class
-
-
 -- | An exact inference transformer that integrates
 -- discrete random variables by enumerating all execution paths.
 newtype Enumerator a = Enumerator (WriterT (Product (Log Double)) [] a)
-  deriving(Functor, Applicative, Monad, Alternative, MonadPlus)
+  deriving (Functor, Applicative, Monad, Alternative, MonadPlus)
 
 instance MonadSample Enumerator where
   random = error "Infinitely supported random variables not supported in Enumerator"
-  bernoulli p = fromList [(True, (Exp . log) p), (False, (Exp . log) (1-p))]
-  categorical v = fromList $ zip [0..] $ map (Exp . log) (V.toList v)
+  bernoulli p = fromList [(True, (Exp . log) p), (False, (Exp . log) (1 - p))]
+  categorical v = fromList $ zip [0 ..] $ map (Exp . log) (V.toList v)
 
 instance MonadCond Enumerator where
   score w = fromList [((), w)]
@@ -60,7 +56,7 @@
 logExplicit (Enumerator m) = map (second getProduct) $ runWriterT m
 
 -- | Same as `toList`, only weights are converted from log-domain.
-explicit :: Enumerator a -> [(a,Double)]
+explicit :: Enumerator a -> [(a, Double)]
 explicit = map (second (exp . ln)) . logExplicit
 
 -- | Returns the model evidence, that is sum of all weights.
@@ -69,13 +65,14 @@
 
 -- | Normalized probability mass of a specific value.
 mass :: Ord a => Enumerator a -> a -> Double
-mass d = f where
-  f a = fromMaybe 0 $ lookup a m
-  m = enumerate d
+mass d = f
+  where
+    f a = fromMaybe 0 $ lookup a m
+    m = enumerate d
 
 -- | Aggregate weights of equal values.
 -- The resulting list is sorted ascendingly according to values.
-compact :: (Num r, Ord a) => [(a,r)] -> [(a,r)]
+compact :: (Num r, Ord a) => [(a, r)] -> [(a, r)]
 compact = Map.toAscList . Map.fromListWith (+)
 
 -- | Aggregate and normalize of weights.
@@ -83,22 +80,25 @@
 --
 -- > enumerate = compact . explicit
 enumerate :: Ord a => Enumerator a -> [(a, Double)]
-enumerate d = compact (zip xs ws) where
-  (xs, ws) = second (map (exp . ln) . normalize) $ unzip (logExplicit d)
+enumerate d = compact (zip xs ws)
+  where
+    (xs, ws) = second (map (exp . ln) . normalize) $ unzip (logExplicit d)
 
 -- | Expectation of a given function computed using normalized weights.
 expectation :: (a -> Double) -> Enumerator a -> Double
 expectation f = Prelude.sum . map (\(x, w) -> f x * (exp . ln) w) . normalizeWeights . logExplicit
 
 normalize :: [Log Double] -> [Log Double]
-normalize xs = map (/ z) xs where
-  z = Log.sum xs
+normalize xs = map (/ z) xs
+  where
+    z = Log.sum xs
 
 -- | Divide all weights by their sum.
 normalizeWeights :: [(a, Log Double)] -> [(a, Log Double)]
-normalizeWeights ls = zip xs ps where
-  (xs, ws) = unzip ls
-  ps = normalize ws
+normalizeWeights ls = zip xs ps
+  where
+    (xs, ws) = unzip ls
+    ps = normalize ws
 
 -- | 'compact' followed by removing values with zero weight.
 normalForm :: Ord a => Enumerator a -> [(a, Double)]
@@ -108,9 +108,11 @@
   p == q = normalForm p == normalForm q
 
 instance Ord a => AEq (Enumerator a) where
-  p === q = xs == ys && ps === qs where
-    (xs,ps) = unzip (normalForm p)
-    (ys,qs) = unzip (normalForm q)
-  p ~== q = xs == ys && ps ~== qs where
-    (xs,ps) = unzip $ filter (not . (~== 0) . snd) $ normalForm p
-    (ys,qs) = unzip $ filter (not . (~== 0) . snd) $ normalForm q
+  p === q = xs == ys && ps === qs
+    where
+      (xs, ps) = unzip (normalForm p)
+      (ys, qs) = unzip (normalForm q)
+  p ~== q = xs == ys && ps ~== qs
+    where
+      (xs, ps) = unzip $ filter (not . (~== 0) . snd) $ normalForm p
+      (ys, qs) = unzip $ filter (not . (~== 0) . snd) $ normalForm q
diff --git a/src/Control/Monad/Bayes/Free.hs b/src/Control/Monad/Bayes/Free.hs
--- a/src/Control/Monad/Bayes/Free.hs
+++ b/src/Control/Monad/Bayes/Free.hs
@@ -1,32 +1,29 @@
-{-|
-Module      : Control.Monad.Bayes.Free
-Description : Free monad transformer over random sampling
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-'FreeSampler' is a free monad transformer over random sampling.
--}
-
-module Control.Monad.Bayes.Free (
-  FreeSampler,
-  hoist,
-  interpret,
-  withRandomness,
-  withPartialRandomness,
-  runWith
-) where
-
-import Data.Functor.Identity
-
-import Control.Monad.Trans
-import Control.Monad.Writer
-import Control.Monad.State
-import Control.Monad.Trans.Free.Church
+-- |
+-- Module      : Control.Monad.Bayes.Free
+-- Description : Free monad transformer over random sampling
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- 'FreeSampler' is a free monad transformer over random sampling.
+module Control.Monad.Bayes.Free
+  ( FreeSampler,
+    hoist,
+    interpret,
+    withRandomness,
+    withPartialRandomness,
+    runWith,
+  )
+where
 
 import Control.Monad.Bayes.Class
+import Control.Monad.State (evalStateT, get, put)
+import Control.Monad.Trans (MonadTrans (..))
+import Control.Monad.Trans.Free.Church (FT, MonadFree (..), hoistFT, iterT, iterTM, liftF)
+import Control.Monad.Writer (WriterT (..), tell)
+import Data.Functor.Identity (Identity, runIdentity)
 
 -- | Random sampling functor.
 newtype SamF a = Random (Double -> a)
@@ -34,15 +31,11 @@
 instance Functor SamF where
   fmap f (Random k) = Random (f . k)
 
-
 -- | Free monad transformer over random sampling.
-
+--
 -- Uses the Church-encoded version of the free monad for efficiency.
-newtype FreeSampler m a = FreeSampler (FT SamF m a)
-  deriving(Functor,Applicative,Monad,MonadTrans)
-
-runFreeSampler :: FreeSampler m a -> FT SamF m a
-runFreeSampler (FreeSampler m) = m
+newtype FreeSampler m a = FreeSampler {runFreeSampler :: FT SamF m a}
+  deriving (Functor, Applicative, Monad, MonadTrans)
 
 instance Monad m => MonadFree SamF (FreeSampler m) where
   wrap = FreeSampler . wrap . fmap runFreeSampler
@@ -56,32 +49,35 @@
 
 -- | Execute random sampling in the transformed monad.
 interpret :: MonadSample m => FreeSampler m a -> m a
-interpret (FreeSampler m) = iterT f m where
-  f (Random k) = random >>= k
+interpret (FreeSampler m) = iterT f m
+  where
+    f (Random k) = random >>= k
 
 -- | Execute computation with supplied values for random choices.
 withRandomness :: Monad m => [Double] -> FreeSampler m a -> m a
-withRandomness randomness (FreeSampler m) = evalStateT (iterTM f m) randomness where
-  f (Random k) = do
-    xs <- get
-    case xs of
-      [] -> error "FreeSampler: the list of randomness was too short"
-      y:ys -> put ys >> k y
+withRandomness randomness (FreeSampler m) = evalStateT (iterTM f m) randomness
+  where
+    f (Random k) = do
+      xs <- get
+      case xs of
+        [] -> error "FreeSampler: the list of randomness was too short"
+        y : ys -> put ys >> k y
 
 -- | Execute computation with supplied values for a subset of random choices.
 -- Return the output value and a record of all random choices used, whether
 -- taken as input or drawn using the transformed monad.
 withPartialRandomness :: MonadSample m => [Double] -> FreeSampler m a -> m (a, [Double])
 withPartialRandomness randomness (FreeSampler m) =
-  runWriterT $ evalStateT (iterTM f $ hoistFT lift m) randomness where
+  runWriterT $ evalStateT (iterTM f $ hoistFT lift m) randomness
+  where
     f (Random k) = do
       -- This block runs in StateT [Double] (WriterT [Double]) m.
       -- StateT propagates consumed randomness while WriterT records
       -- randomness used, whether old or new.
       xs <- get
       x <- case xs of
-            [] -> random
-            y:ys -> put ys >> return y
+        [] -> random
+        y : ys -> put ys >> return y
       tell [x]
       k x
 
diff --git a/src/Control/Monad/Bayes/Helpers.hs b/src/Control/Monad/Bayes/Helpers.hs
--- a/src/Control/Monad/Bayes/Helpers.hs
+++ b/src/Control/Monad/Bayes/Helpers.hs
@@ -1,47 +1,52 @@
-{-|
-Module      : Control.Monad.Bayes.Helpers
-Description : Helper functions for working with inference monads
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Helpers (
-  W,
-  hoistW,
-  P,
-  hoistP,
-  S,
-  hoistS,
-  F,
-  hoistF,
-  T,
-  hoistT,
-  hoistWF,
-  hoistSP,
-  hoistSTP
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Helpers
+-- Description : Helper functions for working with inference monads
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Helpers
+  ( W,
+    hoistW,
+    P,
+    hoistP,
+    S,
+    hoistS,
+    F,
+    hoistF,
+    T,
+    hoistT,
+    hoistWF,
+    hoistSP,
+    hoistSTP,
+  )
+where
 
-import Control.Monad.Bayes.Weighted as Weighted
+import Control.Monad.Bayes.Free as Free
 import Control.Monad.Bayes.Population as Pop
 import Control.Monad.Bayes.Sequential as Seq
-import Control.Monad.Bayes.Free as Free
 import Control.Monad.Bayes.Traced as Tr
+import Control.Monad.Bayes.Weighted as Weighted
 
 type W = Weighted
+
 type P = Population
+
 type S = Sequential
+
 type F = FreeSampler
+
 type T = Traced
 
 hoistW :: (forall x. m x -> n x) -> W m a -> W n a
 hoistW = Weighted.hoist
 
-hoistP :: (Monad m, Monad n)
-      => (forall x. m x -> n x) -> P m a -> P n a
+hoistP ::
+  (Monad m, Monad n) =>
+  (forall x. m x -> n x) ->
+  P m a ->
+  P n a
 hoistP = Pop.hoist
 
 hoistS :: (forall x. m x -> m x) -> S m a -> S m a
@@ -50,18 +55,23 @@
 hoistF :: (Monad m, Monad n) => (forall x. m x -> n x) -> F m a -> F n a
 hoistF = Free.hoist
 
-
-hoistWF :: (Monad m, Monad n)
-      => (forall x. m x -> n x)
-      -> W (F m) a -> W (F n) a
+hoistWF ::
+  (Monad m, Monad n) =>
+  (forall x. m x -> n x) ->
+  W (F m) a ->
+  W (F n) a
 hoistWF m = hoistW $ hoistF m
 
-hoistSP :: Monad m
-        => (forall x. m x -> m x)
-        -> S (P m) a -> S (P m) a
+hoistSP ::
+  Monad m =>
+  (forall x. m x -> m x) ->
+  S (P m) a ->
+  S (P m) a
 hoistSP m = hoistS $ hoistP m
 
-hoistSTP :: Monad m
-         => (forall x. m x -> m x)
-         -> S (T (P m)) a -> S (T (P m)) a
+hoistSTP ::
+  Monad m =>
+  (forall x. m x -> m x) ->
+  S (T (P m)) a ->
+  S (T (P m)) a
 hoistSTP m = hoistS $ hoistT $ hoistP m
diff --git a/src/Control/Monad/Bayes/Inference/PMMH.hs b/src/Control/Monad/Bayes/Inference/PMMH.hs
--- a/src/Control/Monad/Bayes/Inference/PMMH.hs
+++ b/src/Control/Monad/Bayes/Inference/PMMH.hs
@@ -1,38 +1,41 @@
-{-|
-Module      : Control.Monad.Bayes.Inference.PMMH
-Description : Particle Marginal Metropolis-Hastings (PMMH)
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-Particle Marginal Metropolis-Hastings (PMMH) sampling.
-
-Christophe Andrieu, Arnaud Doucet, and Roman Holenstein. 2010. Particle Markov chain Monte Carlo Methods. /Journal of the Royal Statistical Society/ 72 (2010), 269-342. <http://www.stats.ox.ac.uk/~doucet/andrieu_doucet_holenstein_PMCMC.pdf>
--}
-
-module Control.Monad.Bayes.Inference.PMMH (
-  pmmh
-)  where
-
-import Numeric.Log
-
-import Control.Monad.Trans (lift)
+-- |
+-- Module      : Control.Monad.Bayes.Inference.PMMH
+-- Description : Particle Marginal Metropolis-Hastings (PMMH)
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- Particle Marginal Metropolis-Hastings (PMMH) sampling.
+--
+-- Christophe Andrieu, Arnaud Doucet, and Roman Holenstein. 2010. Particle Markov chain Monte Carlo Methods. /Journal of the Royal Statistical Society/ 72 (2010), 269-342. <http://www.stats.ox.ac.uk/~doucet/andrieu_doucet_holenstein_PMCMC.pdf>
+module Control.Monad.Bayes.Inference.PMMH
+  ( pmmh,
+  )
+where
 
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Sequential
+import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Population as Pop
+import Control.Monad.Bayes.Sequential
 import Control.Monad.Bayes.Traced
-import Control.Monad.Bayes.Inference.SMC
+import Control.Monad.Trans (lift)
+import Numeric.Log
 
 -- | Particle Marginal Metropolis-Hastings sampling.
-pmmh :: MonadInfer m
-     => Int -- ^ number of Metropolis-Hastings steps
-     -> Int -- ^ number of time steps
-     -> Int -- ^ number of particles
-     -> Traced m b -- ^ model parameters prior
-     -> (b -> Sequential (Population m) a) -- ^ model
-     -> m [[(a, Log Double)]]
+pmmh ::
+  MonadInfer m =>
+  -- | number of Metropolis-Hastings steps
+  Int ->
+  -- | number of time steps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | model parameters prior
+  Traced m b ->
+  -- | model
+  (b -> Sequential (Population m) a) ->
+  m [[(a, Log Double)]]
 pmmh t k n param model =
   mh t (param >>= runPopulation . pushEvidence . Pop.hoist lift . smcSystematic k n . model)
diff --git a/src/Control/Monad/Bayes/Inference/RMSMC.hs b/src/Control/Monad/Bayes/Inference/RMSMC.hs
--- a/src/Control/Monad/Bayes/Inference/RMSMC.hs
+++ b/src/Control/Monad/Bayes/Inference/RMSMC.hs
@@ -1,69 +1,83 @@
-{-|
-Module      : Control.Monad.Bayes.Inference.RMSMC
-Description : Resample-Move Sequential Monte Carlo (RM-SMC)
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-Resample-move Sequential Monte Carlo (RM-SMC) sampling.
-
-Walter Gilks and Carlo Berzuini. 2001. Following a moving target - Monte Carlo inference for dynamic Bayesian models. /Journal of the Royal Statistical Society/ 63 (2001), 127-146. <http://www.mathcs.emory.edu/~whalen/Papers/BNs/MonteCarlo-DBNs.pdf>
--}
-
-module Control.Monad.Bayes.Inference.RMSMC (
-  rmsmc,
-  rmsmcLocal,
-  rmsmcBasic
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Inference.RMSMC
+-- Description : Resample-Move Sequential Monte Carlo (RM-SMC)
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- Resample-move Sequential Monte Carlo (RM-SMC) sampling.
+--
+-- Walter Gilks and Carlo Berzuini. 2001. Following a moving target - Monte Carlo inference for dynamic Bayesian models. /Journal of the Royal Statistical Society/ 63 (2001), 127-146. <http://www.mathcs.emory.edu/~whalen/Papers/BNs/MonteCarlo-DBNs.pdf>
+module Control.Monad.Bayes.Inference.RMSMC
+  ( rmsmc,
+    rmsmcLocal,
+    rmsmcBasic,
+  )
+where
 
 import Control.Monad.Bayes.Class
+import Control.Monad.Bayes.Helpers
 import Control.Monad.Bayes.Population
 import Control.Monad.Bayes.Sequential as Seq
 import Control.Monad.Bayes.Traced as Tr
-import qualified Control.Monad.Bayes.Traced.Dynamic as TrDyn
 import qualified Control.Monad.Bayes.Traced.Basic as TrBas
-import Control.Monad.Bayes.Helpers
+import qualified Control.Monad.Bayes.Traced.Dynamic as TrDyn
 
 -- | Resample-move Sequential Monte Carlo.
-rmsmc :: MonadSample m
-      => Int -- ^ number of timesteps
-      -> Int -- ^ number of particles
-      -> Int -- ^ number of Metropolis-Hastings transitions after each resampling
-      -> Sequential (Traced (Population m)) a -- ^ model
-      -> Population m a
+rmsmc ::
+  MonadSample m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | number of Metropolis-Hastings transitions after each resampling
+  Int ->
+  -- | model
+  Sequential (Traced (Population m)) a ->
+  Population m a
 rmsmc k n t =
-  marginal .
-  sis (composeCopies t mhStep . hoistT resampleSystematic) k .
-  hoistS (hoistT (spawn n >>))
+  marginal
+    . sis (composeCopies t mhStep . hoistT resampleSystematic) k
+    . hoistS (hoistT (spawn n >>))
 
 -- | Resample-move Sequential Monte Carlo with a more efficient
 -- tracing representation.
-rmsmcBasic :: MonadSample m
-      => Int -- ^ number of timesteps
-      -> Int -- ^ number of particles
-      -> Int -- ^ number of Metropolis-Hastings transitions after each resampling
-      -> Sequential (TrBas.Traced (Population m)) a -- ^ model
-      -> Population m a
+rmsmcBasic ::
+  MonadSample m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | number of Metropolis-Hastings transitions after each resampling
+  Int ->
+  -- | model
+  Sequential (TrBas.Traced (Population m)) a ->
+  Population m a
 rmsmcBasic k n t =
-  TrBas.marginal .
-  sis (composeCopies t TrBas.mhStep . TrBas.hoistT resampleSystematic) k .
-  hoistS (TrBas.hoistT (spawn n >>))
+  TrBas.marginal
+    . sis (composeCopies t TrBas.mhStep . TrBas.hoistT resampleSystematic) k
+    . hoistS (TrBas.hoistT (spawn n >>))
 
 -- | A variant of resample-move Sequential Monte Carlo
 -- where only random variables since last resampling are considered
 -- for rejuvenation.
-rmsmcLocal :: MonadSample m
-           => Int -- ^ number of timesteps
-           -> Int -- ^ number of particles
-           -> Int -- ^ number of Metropolis-Hastings transitions after each resampling
-           -> Sequential (TrDyn.Traced (Population m)) a -- ^ model
-           -> Population m a
+rmsmcLocal ::
+  MonadSample m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | number of Metropolis-Hastings transitions after each resampling
+  Int ->
+  -- | model
+  Sequential (TrDyn.Traced (Population m)) a ->
+  Population m a
 rmsmcLocal k n t =
-  TrDyn.marginal .
-  sis (TrDyn.freeze . composeCopies t TrDyn.mhStep . TrDyn.hoistT resampleSystematic) k .
-  hoistS (TrDyn.hoistT (spawn n >>))
+  TrDyn.marginal
+    . sis (TrDyn.freeze . composeCopies t TrDyn.mhStep . TrDyn.hoistT resampleSystematic) k
+    . hoistS (TrDyn.hoistT (spawn n >>))
 
 -- | Apply a function a given number of times.
 composeCopies :: Int -> (a -> a) -> (a -> a)
diff --git a/src/Control/Monad/Bayes/Inference/SMC.hs b/src/Control/Monad/Bayes/Inference/SMC.hs
--- a/src/Control/Monad/Bayes/Inference/SMC.hs
+++ b/src/Control/Monad/Bayes/Inference/SMC.hs
@@ -1,24 +1,23 @@
-{-|
-Module      : Control.Monad.Bayes.Inference.SMC
-Description : Sequential Monte Carlo (SMC)
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-Sequential Monte Carlo (SMC) sampling.
-
-Arnaud Doucet and Adam M. Johansen. 2011. A tutorial on particle filtering and smoothing: fifteen years later. In /The Oxford Handbook of Nonlinear Filtering/, Dan Crisan and Boris Rozovskii (Eds.). Oxford University Press, Chapter 8.
--}
-
-module Control.Monad.Bayes.Inference.SMC (
-  sir,
-  smcMultinomial,
-  smcSystematic,
-  smcMultinomialPush,
-  smcSystematicPush
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Inference.SMC
+-- Description : Sequential Monte Carlo (SMC)
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- Sequential Monte Carlo (SMC) sampling.
+--
+-- Arnaud Doucet and Adam M. Johansen. 2011. A tutorial on particle filtering and smoothing: fifteen years later. In /The Oxford Handbook of Nonlinear Filtering/, Dan Crisan and Boris Rozovskii (Eds.). Oxford University Press, Chapter 8.
+module Control.Monad.Bayes.Inference.SMC
+  ( sir,
+    smcMultinomial,
+    smcSystematic,
+    smcMultinomialPush,
+    smcSystematicPush,
+  )
+where
 
 import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Population
@@ -26,48 +25,69 @@
 
 -- | Sequential importance resampling.
 -- Basically an SMC template that takes a custom resampler.
-sir :: Monad m
-    => (forall x. Population m x -> Population m x) -- ^ resampler
-    -> Int -- ^ number of timesteps
-    -> Int -- ^ population size
-    -> Sequential (Population m) a -- ^ model
-    -> Population m a
+sir ::
+  Monad m =>
+  -- | resampler
+  (forall x. Population m x -> Population m x) ->
+  -- | number of timesteps
+  Int ->
+  -- | population size
+  Int ->
+  -- | model
+  Sequential (Population m) a ->
+  Population m a
 sir resampler k n = sis resampler k . Seq.hoistFirst (spawn n >>)
 
 -- | Sequential Monte Carlo with multinomial resampling at each timestep.
 -- Weights are not normalized.
-smcMultinomial :: MonadSample m
-               => Int -- ^ number of timesteps
-               -> Int -- ^ number of particles
-               -> Sequential (Population m) a -- ^ model
-               -> Population m a
+smcMultinomial ::
+  MonadSample m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | model
+  Sequential (Population m) a ->
+  Population m a
 smcMultinomial = sir resampleMultinomial
 
 -- | Sequential Monte Carlo with systematic resampling at each timestep.
 -- Weights are not normalized.
-smcSystematic  :: MonadSample m
-               => Int -- ^ number of timesteps
-               -> Int -- ^ number of particles
-               -> Sequential (Population m) a -- ^ model
-               -> Population m a
+smcSystematic ::
+  MonadSample m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | model
+  Sequential (Population m) a ->
+  Population m a
 smcSystematic = sir resampleSystematic
 
 -- | Sequential Monte Carlo with multinomial resampling at each timestep.
 -- Weights are normalized at each timestep and the total weight is pushed
 -- as a score into the transformed monad.
-smcMultinomialPush :: MonadInfer m
-                   => Int -- ^ number of timesteps
-                   -> Int -- ^ number of particles
-                   -> Sequential (Population m) a -- ^ model
-                   -> Population m a
+smcMultinomialPush ::
+  MonadInfer m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | model
+  Sequential (Population m) a ->
+  Population m a
 smcMultinomialPush = sir (pushEvidence . resampleMultinomial)
 
 -- | Sequential Monte Carlo with systematic resampling at each timestep.
 -- Weights are normalized at each timestep and the total weight is pushed
 -- as a score into the transformed monad.
-smcSystematicPush  :: MonadInfer m
-                   => Int -- ^ number of timesteps
-                   -> Int -- ^ number of particles
-                   -> Sequential (Population m) a -- ^ model
-                   -> Population m a
+smcSystematicPush ::
+  MonadInfer m =>
+  -- | number of timesteps
+  Int ->
+  -- | number of particles
+  Int ->
+  -- | model
+  Sequential (Population m) a ->
+  Population m a
 smcSystematicPush = sir (pushEvidence . resampleSystematic)
diff --git a/src/Control/Monad/Bayes/Inference/SMC2.hs b/src/Control/Monad/Bayes/Inference/SMC2.hs
--- a/src/Control/Monad/Bayes/Inference/SMC2.hs
+++ b/src/Control/Monad/Bayes/Inference/SMC2.hs
@@ -1,33 +1,32 @@
-{-|
-Module      : Control.Monad.Bayes.Inference.SMC2
-Description : Sequential Monte Carlo squared (SMC²)
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-Sequential Monte Carlo squared (SMC²) sampling.
-
-Nicolas Chopin, Pierre E. Jacob, and Omiros Papaspiliopoulos. 2013. SMC²: an efficient algorithm for sequential analysis of state space models. /Journal of the Royal Statistical Society Series B: Statistical Methodology/ 75 (2013), 397-426. Issue 3. <https://doi.org/10.1111/j.1467-9868.2012.01046.x>
--}
-
-module Control.Monad.Bayes.Inference.SMC2 (
-  smc2
-) where
-
-import Numeric.Log
-import Control.Monad.Trans
+-- |
+-- Module      : Control.Monad.Bayes.Inference.SMC2
+-- Description : Sequential Monte Carlo squared (SMC²)
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- Sequential Monte Carlo squared (SMC²) sampling.
+--
+-- Nicolas Chopin, Pierre E. Jacob, and Omiros Papaspiliopoulos. 2013. SMC²: an efficient algorithm for sequential analysis of state space models. /Journal of the Royal Statistical Society Series B: Statistical Methodology/ 75 (2013), 397-426. Issue 3. <https://doi.org/10.1111/j.1467-9868.2012.01046.x>
+module Control.Monad.Bayes.Inference.SMC2
+  ( smc2,
+  )
+where
 
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Population as Pop
-import Control.Monad.Bayes.Inference.SMC
-import Control.Monad.Bayes.Inference.RMSMC
 import Control.Monad.Bayes.Helpers
+import Control.Monad.Bayes.Inference.RMSMC
+import Control.Monad.Bayes.Inference.SMC
+import Control.Monad.Bayes.Population as Pop
+import Control.Monad.Trans
+import Numeric.Log
 
 -- | Helper monad transformer for preprocessing the model for 'smc2'.
 newtype SMC2 m a = SMC2 (S (T (P m)) a)
-  deriving(Functor,Applicative,Monad)
+  deriving (Functor, Applicative, Monad)
+
 setup :: SMC2 m a -> S (T (P m)) a
 setup (SMC2 m) = m
 
@@ -43,13 +42,20 @@
 instance MonadSample m => MonadInfer (SMC2 m)
 
 -- | Sequential Monte Carlo squared.
-smc2 :: MonadSample m
-     => Int -- ^ number of time steps
-     -> Int -- ^ number of inner particles
-     -> Int -- ^ number of outer particles
-     -> Int -- ^ number of MH transitions
-     -> S (T (P m)) b -- ^ model parameters
-     -> ( b -> S (P (SMC2 m)) a) -- ^ model
-     -> P m [(a, Log Double)]
+smc2 ::
+  MonadSample m =>
+  -- | number of time steps
+  Int ->
+  -- | number of inner particles
+  Int ->
+  -- | number of outer particles
+  Int ->
+  -- | number of MH transitions
+  Int ->
+  -- | model parameters
+  S (T (P m)) b ->
+  -- | model
+  (b -> S (P (SMC2 m)) a) ->
+  P m [(a, Log Double)]
 smc2 k n p t param model =
   rmsmc k p t (param >>= setup . runPopulation . smcSystematicPush k n . model)
diff --git a/src/Control/Monad/Bayes/Population.hs b/src/Control/Monad/Bayes/Population.hs
--- a/src/Control/Monad/Bayes/Population.hs
+++ b/src/Control/Monad/Bayes/Population.hs
@@ -1,17 +1,15 @@
-{-|
-Module      : Control.Monad.Bayes.Population
-Description : Representation of distributions using multiple samples
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-'Population' turns a single sample into a collection of weighted samples.
--}
-
-module Control.Monad.Bayes.Population (
-    Population,
+-- |
+-- Module      : Control.Monad.Bayes.Population
+-- Description : Representation of distributions using multiple samples
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- 'Population' turns a single sample into a collection of weighted samples.
+module Control.Monad.Bayes.Population
+  ( Population,
     runPopulation,
     explicitPopulation,
     fromWeightedList,
@@ -27,26 +25,24 @@
     normalize,
     popAvg,
     flatten,
-    hoist
-                 ) where
-
-import Prelude hiding (sum, all)
+    hoist,
+  )
+where
 
 import Control.Arrow (second)
+import Control.Monad (replicateM)
+import Control.Monad.Bayes.Class
+import Control.Monad.Bayes.Weighted hiding (flatten, hoist)
 import Control.Monad.Trans
 import Control.Monad.Trans.List
-import Control.Monad (replicateM)
-
 import qualified Data.List
 import qualified Data.Vector as V
-
 import Numeric.Log
-import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Weighted hiding (flatten, hoist)
+import Prelude hiding (all, sum)
 
 -- | A collection of weighted samples, or particles.
 newtype Population m a = Population (Weighted (ListT m) a)
-  deriving(Functor,Applicative,Monad,MonadIO,MonadSample,MonadCond,MonadInfer)
+  deriving (Functor, Applicative, Monad, MonadIO, MonadSample, MonadCond, MonadInfer)
 
 instance MonadTrans Population where
   lift = Population . lift . lift
@@ -61,7 +57,7 @@
 explicitPopulation = fmap (map (second (exp . ln))) . runPopulation
 
 -- | Initialize 'Population' with a concrete weighted sample.
-fromWeightedList :: Monad m => m [(a,Log Double)] -> Population m a
+fromWeightedList :: Monad m => m [(a, Log Double)] -> Population m a
 fromWeightedList = Population . withWeight . ListT
 
 -- | Increase the sample size by a given factor.
@@ -71,41 +67,46 @@
 spawn :: Monad m => Int -> Population m ()
 spawn n = fromWeightedList $ pure $ replicate n ((), 1 / fromIntegral n)
 
-resampleGeneric :: MonadSample m
-                => (V.Vector Double -> m [Int]) -- ^ resampler
-                -> Population m a -> Population m a
+resampleGeneric ::
+  MonadSample m =>
+  -- | resampler
+  (V.Vector Double -> m [Int]) ->
+  Population m a ->
+  Population m a
 resampleGeneric resampler m = fromWeightedList $ do
   pop <- runPopulation m
   let (xs, ps) = unzip pop
   let n = length xs
   let z = sum ps
-  if z > 0 then do
-    let weights = V.fromList (map (exp . ln . (/z)) ps)
-    ancestors <- resampler weights
-    let xvec = V.fromList xs
-    let offsprings = map (xvec V.!) ancestors
-    return $ map (, z / fromIntegral n) offsprings
-  else
-    -- if all weights are zero do not resample
-    return pop
+  if z > 0
+    then do
+      let weights = V.fromList (map (exp . ln . (/ z)) ps)
+      ancestors <- resampler weights
+      let xvec = V.fromList xs
+      let offsprings = map (xvec V.!) ancestors
+      return $ map (,z / fromIntegral n) offsprings
+    else-- if all weights are zero do not resample
+      return pop
 
 -- | Systematic resampling helper.
 systematic :: Double -> V.Vector Double -> [Int]
-systematic u ps = f 0 (u / fromIntegral n) 0 0 [] where
-  prob i = ps V.! i
-  n = length ps
-  inc = 1 / fromIntegral n
-  f i _ _ _ acc | i == n = acc
-  f i v j q acc =
-    if v < q then
-      f (i+1) (v+inc) j q (j-1:acc)
-    else
-      f i v (j + 1) (q + prob j) acc
+systematic u ps = f 0 (u / fromIntegral n) 0 0 []
+  where
+    prob i = ps V.! i
+    n = length ps
+    inc = 1 / fromIntegral n
+    f i _ _ _ acc | i == n = acc
+    f i v j q acc =
+      if v < q
+        then f (i + 1) (v + inc) j q (j -1 : acc)
+        else f i v (j + 1) (q + prob j) acc
 
 -- | Resample the population using the underlying monad and a systematic resampling scheme.
 -- The total weight is preserved.
-resampleSystematic :: (MonadSample m)
-         => Population m a -> Population m a
+resampleSystematic ::
+  (MonadSample m) =>
+  Population m a ->
+  Population m a
 resampleSystematic = resampleGeneric (\ps -> (`systematic` ps) <$> random)
 
 -- | Multinomial resampler.
@@ -114,14 +115,18 @@
 
 -- | Resample the population using the underlying monad and a multinomial resampling scheme.
 -- The total weight is preserved.
-resampleMultinomial :: (MonadSample m)
-         => Population m a -> Population m a
+resampleMultinomial ::
+  (MonadSample m) =>
+  Population m a ->
+  Population m a
 resampleMultinomial = resampleGeneric multinomial
 
 -- | Separate the sum of weights into the 'Weighted' transformer.
 -- Weights are normalized after this operation.
-extractEvidence :: Monad m
-                => Population m a -> Population (Weighted m) a
+extractEvidence ::
+  Monad m =>
+  Population m a ->
+  Population (Weighted m) a
 extractEvidence m = fromWeightedList $ do
   pop <- lift $ runPopulation m
   let (xs, ps) = unzip pop
@@ -132,14 +137,18 @@
 
 -- | Push the evidence estimator as a score to the transformed monad.
 -- Weights are normalized after this operation.
-pushEvidence :: MonadCond m
-           => Population m a -> Population m a
+pushEvidence ::
+  MonadCond m =>
+  Population m a ->
+  Population m a
 pushEvidence = hoist applyWeight . extractEvidence
 
 -- | A properly weighted single sample, that is one picked at random according
 -- to the weights, with the sum of all weights.
-proper :: (MonadSample m)
-       => Population m a -> Weighted m a
+proper ::
+  (MonadSample m) =>
+  Population m a ->
+  Weighted m a
 proper m = do
   pop <- runPopulation $ extractEvidence m
   let (xs, ps) = unzip pop
@@ -155,13 +164,18 @@
 -- in the transformed monad.
 -- This way a single sample can be selected from a population without
 -- introducing bias.
-collapse :: (MonadInfer m)
-         => Population m a -> m a
+collapse ::
+  (MonadInfer m) =>
+  Population m a ->
+  m a
 collapse = applyWeight . proper
 
 -- | Applies a random transformation to a population.
-mapPopulation :: (Monad m) => ([(a, Log Double)] -> m [(a, Log Double)]) ->
-  Population m a -> Population m a
+mapPopulation ::
+  (Monad m) =>
+  ([(a, Log Double)] -> m [(a, Log Double)]) ->
+  Population m a ->
+  Population m a
 mapPopulation f m = fromWeightedList $ runPopulation m >>= f
 
 -- | Normalizes the weights in the population so that their sum is 1.
@@ -173,20 +187,24 @@
 popAvg :: (Monad m) => (a -> Double) -> Population m a -> m Double
 popAvg f p = do
   xs <- explicitPopulation p
-  let ys = map (\(x,w) -> f x * w) xs
+  let ys = map (\(x, w) -> f x * w) xs
   let t = Data.List.sum ys
   return t
 
 -- | Combine a population of populations into a single population.
 flatten :: Monad m => Population (Population m) a -> Population m a
-flatten m = Population $ withWeight $ ListT t where
-  t = f <$> (runPopulation . runPopulation) m
-  f d = do
-    (x,p) <- d
-    (y,q) <- x
-    return (y, p*q)
+flatten m = Population $ withWeight $ ListT t
+  where
+    t = f <$> (runPopulation . runPopulation) m
+    f d = do
+      (x, p) <- d
+      (y, q) <- x
+      return (y, p * q)
 
 -- | Applies a transformation to the inner monad.
-hoist :: (Monad m, Monad n)
-      => (forall x. m x -> n x) -> Population m a -> Population n a
+hoist ::
+  (Monad m, Monad n) =>
+  (forall x. m x -> n x) ->
+  Population m a ->
+  Population n a
 hoist f = fromWeightedList . f . runPopulation
diff --git a/src/Control/Monad/Bayes/Sampler.hs b/src/Control/Monad/Bayes/Sampler.hs
--- a/src/Control/Monad/Bayes/Sampler.hs
+++ b/src/Control/Monad/Bayes/Sampler.hs
@@ -1,40 +1,38 @@
-{-|
-Module      : Control.Monad.Bayes.Sampler
-Description : Pseudo-random sampling monads
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-'SamplerIO' and 'SamplerST' are instances of 'MonadSample'. Apply a 'MonadCond'
-transformer to obtain a 'MonadInfer' that can execute probabilistic models.
--}
-
-module Control.Monad.Bayes.Sampler (
-    SamplerIO,
+-- |
+-- Module      : Control.Monad.Bayes.Sampler
+-- Description : Pseudo-random sampling monads
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- 'SamplerIO' and 'SamplerST' are instances of 'MonadSample'. Apply a 'MonadCond'
+-- transformer to obtain a 'MonadInfer' that can execute probabilistic models.
+module Control.Monad.Bayes.Sampler
+  ( SamplerIO,
     sampleIO,
     sampleIOfixed,
     sampleIOwith,
     Seed,
-    SamplerST(SamplerST),
+    SamplerST (SamplerST),
     runSamplerST,
     sampleST,
-    sampleSTfixed
-               ) where
+    sampleSTfixed,
+  )
+where
 
+import Control.Monad.Bayes.Class
 import Control.Monad.ST (ST, runST, stToIO)
+import Control.Monad.State (State, state)
+import Control.Monad.Trans (MonadIO, lift)
+import Control.Monad.Trans.Reader (ReaderT, ask, mapReaderT, runReaderT)
 import System.Random.MWC
 import qualified System.Random.MWC.Distributions as MWC
-import Control.Monad.State (State, state)
-import Control.Monad.Trans (lift, MonadIO)
-import Control.Monad.Trans.Reader (ReaderT, runReaderT, ask, mapReaderT)
 
-import Control.Monad.Bayes.Class
-
 -- | An 'IO' based random sampler using the MWC-Random package.
 newtype SamplerIO a = SamplerIO (ReaderT GenIO IO a)
-  deriving(Functor, Applicative, Monad, MonadIO)
+  deriving (Functor, Applicative, Monad, MonadIO)
 
 -- | Initialize a pseudo-random number generator using randomness supplied by
 -- the operating system.
@@ -58,9 +56,6 @@
 instance MonadSample SamplerIO where
   random = fromSamplerST random
 
-
-
-
 -- | An 'ST' based random sampler using the @mwc-random@ package.
 newtype SamplerST a = SamplerST (forall s. ReaderT (GenST s) (ST s) a)
 
@@ -100,7 +95,7 @@
 instance MonadSample SamplerST where
   random = fromMWC System.Random.MWC.uniform
 
-  uniform a b = fromMWC $ uniformR (a,b)
+  uniform a b = fromMWC $ uniformR (a, b)
   normal m s = fromMWC $ MWC.normal m s
   gamma shape scale = fromMWC $ MWC.gamma shape scale
   beta a b = fromMWC $ MWC.beta a b
diff --git a/src/Control/Monad/Bayes/Sequential.hs b/src/Control/Monad/Bayes/Sequential.hs
--- a/src/Control/Monad/Bayes/Sequential.hs
+++ b/src/Control/Monad/Bayes/Sequential.hs
@@ -1,40 +1,38 @@
-{-|
-Module      : Control.Monad.Bayes.Sequential
-Description : Suspendable probabilistic computation
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-'Sequential' represents a computation that can be suspended.
--}
-
-module Control.Monad.Bayes.Sequential (
-    Sequential,
+-- |
+-- Module      : Control.Monad.Bayes.Sequential
+-- Description : Suspendable probabilistic computation
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- 'Sequential' represents a computation that can be suspended.
+module Control.Monad.Bayes.Sequential
+  ( Sequential,
     suspend,
     finish,
     advance,
     finished,
     hoistFirst,
     hoist,
-    sis
-                ) where
+    sis,
+  )
+where
 
-import Control.Monad.Trans
+import Control.Monad.Bayes.Class
 import Control.Monad.Coroutine hiding (suspend)
 import Control.Monad.Coroutine.SuspensionFunctors
+import Control.Monad.Trans
 import Data.Either
 
-import Control.Monad.Bayes.Class
-
 -- | Represents a computation that can be suspended at certain points.
 -- The intermediate monadic effects can be extracted, which is particularly
 -- useful for implementation of Sequential Monte Carlo related methods.
 -- All the probabilistic effects are lifted from the transformed monad, but
 -- also `suspend` is inserted after each `factor`.
 newtype Sequential m a = Sequential {runSequential :: Coroutine (Await ()) m a}
-  deriving(Functor,Applicative,Monad,MonadTrans,MonadIO)
+  deriving (Functor, Applicative, Monad, MonadTrans, MonadIO)
 
 extract :: Await () a -> a
 extract (Await f) = f ()
@@ -76,8 +74,11 @@
 
 -- | Transform the inner monad.
 -- The transformation is applied recursively through all the suspension points.
-hoist :: (Monad m, Monad n) =>
-            (forall x. m x -> n x) -> Sequential m a -> Sequential n a
+hoist ::
+  (Monad m, Monad n) =>
+  (forall x. m x -> n x) ->
+  Sequential m a ->
+  Sequential n a
 hoist f = Sequential . mapMonad f . runSequential
 
 -- | Apply a function a given number of times.
@@ -86,9 +87,12 @@
 
 -- | Sequential importance sampling.
 -- Applies a given transformation after each time step.
-sis :: Monad m
-    => (forall x. m x -> m x) -- ^ transformation
-    -> Int -- ^ number of time steps
-    -> Sequential m a
-    -> m a
+sis ::
+  Monad m =>
+  -- | transformation
+  (forall x. m x -> m x) ->
+  -- | number of time steps
+  Int ->
+  Sequential m a ->
+  m a
 sis f k = finish . composeCopies k (advance . hoistFirst f)
diff --git a/src/Control/Monad/Bayes/Traced.hs b/src/Control/Monad/Bayes/Traced.hs
--- a/src/Control/Monad/Bayes/Traced.hs
+++ b/src/Control/Monad/Bayes/Traced.hs
@@ -1,16 +1,14 @@
-{-|
-Module      : Control.Monad.Bayes.Traced
-Description : Distributions on execution traces
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Traced (
-  module Control.Monad.Bayes.Traced.Static
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Traced
+-- Description : Distributions on execution traces
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Traced
+  ( module Control.Monad.Bayes.Traced.Static,
+  )
+where
 
 import Control.Monad.Bayes.Traced.Static
diff --git a/src/Control/Monad/Bayes/Traced/Basic.hs b/src/Control/Monad/Bayes/Traced/Basic.hs
--- a/src/Control/Monad/Bayes/Traced/Basic.hs
+++ b/src/Control/Monad/Bayes/Traced/Basic.hs
@@ -1,41 +1,35 @@
-{-|
-Module      : Control.Monad.Bayes.Traced.Basic
-Description : Distributions on full execution traces of full programs
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Traced.Basic (
-  Traced,
-  hoistT,
-  marginal,
-  mhStep,
-  mh
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Traced.Basic
+-- Description : Distributions on full execution traces of full programs
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Traced.Basic
+  ( Traced,
+    hoistT,
+    marginal,
+    mhStep,
+    mh,
+  )
+where
 
-import Data.Functor.Identity
 import Control.Applicative (liftA2)
-
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Weighted as Weighted
-import Control.Monad.Bayes.Free as FreeSampler
-
+import Control.Monad.Bayes.Free (FreeSampler)
 import Control.Monad.Bayes.Traced.Common
+import Control.Monad.Bayes.Weighted (Weighted)
+import Data.Functor.Identity (Identity)
 
 -- | Tracing monad that records random choices made in the program.
--- The first component is used to run the program with a modified trace,
--- while the second records a trace and an output value from a run.
-data Traced m a = Traced (Weighted (FreeSampler Identity) a) (m (Trace a))
-
-traceDist :: Traced m a -> m (Trace a)
-traceDist (Traced _ d) = d
-
-model :: Traced m a -> Weighted (FreeSampler Identity) a
-model (Traced m _) = m
+data Traced m a
+  = Traced
+      { -- | Run the program with a modified trace.
+        model :: Weighted (FreeSampler Identity) a,
+        -- | Record trace and output.
+        traceDist :: m (Trace a)
+      }
 
 instance Monad m => Functor (Traced m) where
   fmap f (Traced m d) = Traced (fmap f m) (fmap (fmap f) d)
@@ -45,9 +39,10 @@
   (Traced mf df) <*> (Traced mx dx) = Traced (mf <*> mx) (liftA2 (<*>) df dx)
 
 instance Monad m => Monad (Traced m) where
-  (Traced mx dx) >>= f = Traced my dy where
-    my = mx >>= model . f
-    dy = dx `bind` (traceDist . f)
+  (Traced mx dx) >>= f = Traced my dy
+    where
+      my = mx >>= model . f
+      dy = dx `bind` (traceDist . f)
 
 instance MonadSample m => MonadSample (Traced m) where
   random = Traced random (fmap singleton random)
@@ -64,18 +59,19 @@
 marginal :: Monad m => Traced m a -> m a
 marginal (Traced _ d) = fmap output d
 
--- | A single step of the Trace MH algorithm.
+-- | A single step of the Trace Metropolis-Hastings algorithm.
 mhStep :: MonadSample m => Traced m a -> Traced m a
-mhStep (Traced m d) = Traced m d' where
-  d' = d >>= mhTrans' m
+mhStep (Traced m d) = Traced m d'
+  where
+    d' = d >>= mhTrans' m
 
--- | Full run of the Trace MH algorithm with a specified
+-- | Full run of the Trace Metropolis-Hastings algorithm with a specified
 -- number of steps.
 mh :: MonadSample m => Int -> Traced m a -> m [a]
-mh n (Traced m d) = fmap (map output) t where
-  t = f n
-  f 0 = fmap (:[]) d
-  f k = do
-    ~(x:xs) <- f (k-1)
-    y <- mhTrans' m x
-    return (y:x:xs)
+mh n (Traced m d) = fmap (map output) (f n)
+  where
+    f 0 = fmap (: []) d
+    f k = do
+      ~(x : xs) <- f (k -1)
+      y <- mhTrans' m x
+      return (y : x : xs)
diff --git a/src/Control/Monad/Bayes/Traced/Common.hs b/src/Control/Monad/Bayes/Traced/Common.hs
--- a/src/Control/Monad/Bayes/Traced/Common.hs
+++ b/src/Control/Monad/Bayes/Traced/Common.hs
@@ -1,52 +1,57 @@
-{-|
-Module      : Control.Monad.Bayes.Traced.Common
-Description : Numeric code for Trace MCMC
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Traced.Common (
-  Trace,
-  singleton,
-  output,
-  scored,
-  bind,
-  mhTrans,
-  mhTrans'
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Traced.Common
+-- Description : Numeric code for Trace MCMC
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Traced.Common
+  ( Trace,
+    singleton,
+    output,
+    scored,
+    bind,
+    mhTrans,
+    mhTrans',
+  )
+where
 
+import Control.Monad.Bayes.Class
+import Control.Monad.Bayes.Free as FreeSampler
+import Control.Monad.Bayes.Weighted as Weighted
 import Control.Monad.Trans.Writer
-import qualified Data.Vector as V
 import Data.Functor.Identity
-
 import Numeric.Log (Log, ln)
-
-import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Weighted as Weighted
-import Control.Monad.Bayes.Free as FreeSampler
+import Statistics.Distribution.DiscreteUniform (discreteUniformAB)
 
-data Trace a =
-  Trace {
-    variables :: [Double],
-    output :: a,
-    density :: Log Double
-  }
+-- | Collection of random variables sampled during the program's execution.
+data Trace a
+  = Trace
+      { -- | Sequence of random variables sampled during the program's execution.
+        variables :: [Double],
+        -- |
+        output :: a,
+        -- | The probability of observing this particular sequence.
+        density :: Log Double
+      }
 
 instance Functor Trace where
   fmap f t = t {output = f (output t)}
 
 instance Applicative Trace where
   pure x = Trace {variables = [], output = x, density = 1}
-  tf <*> tx = Trace {variables = variables tf ++ variables tx, output = output tf (output tx), density = density tf * density tx}
+  tf <*> tx =
+    Trace
+      { variables = variables tf ++ variables tx,
+        output = output tf (output tx),
+        density = density tf * density tx
+      }
 
 instance Monad Trace where
   t >>= f =
-    let t' = f (output t) in
-    t' {variables = variables t ++ variables t', density = density t * density t'}
+    let t' = f (output t)
+     in t' {variables = variables t ++ variables t', density = density t * density t'}
 
 singleton :: Double -> Trace Double
 singleton u = Trace {variables = [u], output = u, density = 1}
@@ -62,17 +67,15 @@
 
 -- | A single Metropolis-corrected transition of single-site Trace MCMC.
 mhTrans :: MonadSample m => Weighted (FreeSampler m) a -> Trace a -> m (Trace a)
-mhTrans m t = do
-  let us = variables t
-      p = density t
+mhTrans m t@Trace {variables = us, density = p} = do
+  let n = length us
   us' <- do
-    let n = length us
-    i <- categorical $ V.replicate n (1 / fromIntegral n)
+    i <- discrete $ discreteUniformAB 0 (n -1)
     u' <- random
-    let (xs, _:ys) = splitAt i us
-    return $ xs ++ (u':ys)
+    let (xs, _ : ys) = splitAt i us
+    return $ xs ++ (u' : ys)
   ((b, q), vs) <- runWriterT $ runWeighted $ Weighted.hoist (WriterT . withPartialRandomness us') m
-  let ratio = (exp . ln) $ min 1 (q * fromIntegral (length us) / (p * fromIntegral (length vs)))
+  let ratio = (exp . ln) $ min 1 (q * fromIntegral n / (p * fromIntegral (length vs)))
   accept <- bernoulli ratio
   return $ if accept then Trace vs b q else t
 
diff --git a/src/Control/Monad/Bayes/Traced/Dynamic.hs b/src/Control/Monad/Bayes/Traced/Dynamic.hs
--- a/src/Control/Monad/Bayes/Traced/Dynamic.hs
+++ b/src/Control/Monad/Bayes/Traced/Dynamic.hs
@@ -1,37 +1,31 @@
-{-|
-Module      : Control.Monad.Bayes.Traced.Dynamic
-Description : Distributions on execution traces that can be dynamically frozen
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Traced.Dynamic (
-  Traced,
-  hoistT,
-  marginal,
-  freeze,
-  mhStep,
-  mh
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Traced.Dynamic
+-- Description : Distributions on execution traces that can be dynamically frozen
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Traced.Dynamic
+  ( Traced,
+    hoistT,
+    marginal,
+    freeze,
+    mhStep,
+    mh,
+  )
+where
 
 import Control.Monad (join)
-import Control.Monad.Trans
-
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Weighted as Weighted
-import Control.Monad.Bayes.Free as FreeSampler
-
+import Control.Monad.Bayes.Free (FreeSampler)
 import Control.Monad.Bayes.Traced.Common
+import Control.Monad.Bayes.Weighted (Weighted)
+import Control.Monad.Trans (MonadTrans (..))
 
--- | A tracing monad where only a subset of random choices are traced
--- and this subset can be adjusted dynamically.
-newtype Traced m a = Traced (m (Weighted (FreeSampler m) a, Trace a))
-runTraced :: Traced m a -> m (Weighted (FreeSampler m) a, Trace a)
-runTraced (Traced c) = c
+-- | A tracing monad where only a subset of random choices are traced and this
+-- subset can be adjusted dynamically.
+newtype Traced m a = Traced {runTraced :: m (Weighted (FreeSampler m) a, Trace a)}
 
 pushM :: Monad m => m (Weighted (FreeSampler m) a) -> Weighted (FreeSampler m) a
 pushM = join . lift . lift
@@ -71,31 +65,35 @@
 hoistT :: (forall x. m x -> m x) -> Traced m a -> Traced m a
 hoistT f (Traced c) = Traced (f c)
 
+-- | Discard the trace and supporting infrastructure.
 marginal :: Monad m => Traced m a -> m a
 marginal (Traced c) = fmap (output . snd) c
 
--- | Freeze all traced random choices to their current
--- values and stop tracing them.
+-- | Freeze all traced random choices to their current values and stop tracing
+-- them.
 freeze :: Monad m => Traced m a -> Traced m a
 freeze (Traced c) = Traced $ do
   (_, t) <- c
   let x = output t
   return (return x, pure x)
 
+-- | A single step of the Trace Metropolis-Hastings algorithm.
 mhStep :: MonadSample m => Traced m a -> Traced m a
 mhStep (Traced c) = Traced $ do
   (m, t) <- c
   t' <- mhTrans m t
   return (m, t')
 
+-- | Full run of the Trace Metropolis-Hastings algorithm with a specified
+-- number of steps.
 mh :: MonadSample m => Int -> Traced m a -> m [a]
 mh n (Traced c) = do
-  (m,t) <- c
+  (m, t) <- c
   let f 0 = return [t]
       f k = do
-        ~(x:xs) <- f (k-1)
+        ~(x : xs) <- f (k -1)
         y <- mhTrans m x
-        return (y:x:xs)
+        return (y : x : xs)
   ts <- f n
   let xs = map output ts
   return xs
diff --git a/src/Control/Monad/Bayes/Traced/Static.hs b/src/Control/Monad/Bayes/Traced/Static.hs
--- a/src/Control/Monad/Bayes/Traced/Static.hs
+++ b/src/Control/Monad/Bayes/Traced/Static.hs
@@ -1,41 +1,36 @@
-{-|
-Module      : Control.Monad.Bayes.Traced.Static
-Description : Distributions on execution traces of full programs
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
--}
-
-module Control.Monad.Bayes.Traced.Static (
-  Traced,
-  hoistT,
-  marginal,
-  mhStep,
-  mh
-) where
+-- |
+-- Module      : Control.Monad.Bayes.Traced.Static
+-- Description : Distributions on execution traces of full programs
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+module Control.Monad.Bayes.Traced.Static
+  ( Traced,
+    hoistT,
+    marginal,
+    mhStep,
+    mh,
+  )
+where
 
-import Control.Monad.Trans
 import Control.Applicative (liftA2)
-
 import Control.Monad.Bayes.Class
-import Control.Monad.Bayes.Weighted as Weighted
-import Control.Monad.Bayes.Free as FreeSampler
-
+import Control.Monad.Bayes.Free (FreeSampler)
 import Control.Monad.Bayes.Traced.Common
+import Control.Monad.Bayes.Weighted (Weighted)
+import Control.Monad.Trans (MonadTrans (..))
 
 -- | A tracing monad where only a subset of random choices are traced.
--- The random choices that are not to be traced should be lifted
--- from the transformed monad.
-data Traced m a = Traced (Weighted (FreeSampler m) a) (m (Trace a))
-
-traceDist :: Traced m a -> m (Trace a)
-traceDist (Traced _ d) = d
-
-model :: Traced m a -> Weighted (FreeSampler m) a
-model (Traced m _) = m
+--
+-- The random choices that are not to be traced should be lifted from the
+-- transformed monad.
+data Traced m a
+  = Traced
+      { model :: Weighted (FreeSampler m) a,
+        traceDist :: m (Trace a)
+      }
 
 instance Monad m => Functor (Traced m) where
   fmap f (Traced m d) = Traced (fmap f m) (fmap (fmap f) d)
@@ -45,9 +40,10 @@
   (Traced mf df) <*> (Traced mx dx) = Traced (mf <*> mx) (liftA2 (<*>) df dx)
 
 instance Monad m => Monad (Traced m) where
-  (Traced mx dx) >>= f = Traced my dy where
-    my = mx >>= model . f
-    dy = dx `bind` (traceDist . f)
+  (Traced mx dx) >>= f = Traced my dy
+    where
+      my = mx >>= model . f
+      dy = dx `bind` (traceDist . f)
 
 instance MonadTrans Traced where
   lift m = Traced (lift $ lift m) (fmap pure m)
@@ -63,18 +59,23 @@
 hoistT :: (forall x. m x -> m x) -> Traced m a -> Traced m a
 hoistT f (Traced m d) = Traced m (f d)
 
+-- | Discard the trace and supporting infrastructure.
 marginal :: Monad m => Traced m a -> m a
 marginal (Traced _ d) = fmap output d
 
+-- | A single step of the Trace Metropolis-Hastings algorithm.
 mhStep :: MonadSample m => Traced m a -> Traced m a
-mhStep (Traced m d) = Traced m d' where
-  d' = d >>= mhTrans m
+mhStep (Traced m d) = Traced m d'
+  where
+    d' = d >>= mhTrans m
 
+-- | Full run of the Trace Metropolis-Hastings algorithm with a specified
+-- number of steps.
 mh :: MonadSample m => Int -> Traced m a -> m [a]
-mh n (Traced m d) = fmap (map output) t where
-  t = f n
-  f 0 = fmap (:[]) d
-  f k = do
-    ~(x:xs) <- f (k-1)
-    y <- mhTrans m x
-    return (y:x:xs)
+mh n (Traced m d) = fmap (map output) (f n)
+  where
+    f 0 = fmap (: []) d
+    f k = do
+      ~(x : xs) <- f (k -1)
+      y <- mhTrans m x
+      return (y : x : xs)
diff --git a/src/Control/Monad/Bayes/Weighted.hs b/src/Control/Monad/Bayes/Weighted.hs
--- a/src/Control/Monad/Bayes/Weighted.hs
+++ b/src/Control/Monad/Bayes/Weighted.hs
@@ -1,18 +1,16 @@
-{-|
-Module      : Control.Monad.Bayes.Weighted
-Description : Probability monad accumulating the likelihood
-Copyright   : (c) Adam Scibior, 2015-2020
-License     : MIT
-Maintainer  : leonhard.markert@tweag.io
-Stability   : experimental
-Portability : GHC
-
-'Weighted' is an instance of 'MonadCond'. Apply a 'MonadSample' transformer to
-obtain a 'MonadInfer' that can execute probabilistic models.
--}
-
-module Control.Monad.Bayes.Weighted (
-    Weighted,
+-- |
+-- Module      : Control.Monad.Bayes.Weighted
+-- Description : Probability monad accumulating the likelihood
+-- Copyright   : (c) Adam Scibior, 2015-2020
+-- License     : MIT
+-- Maintainer  : leonhard.markert@tweag.io
+-- Stability   : experimental
+-- Portability : GHC
+--
+-- 'Weighted' is an instance of 'MonadCond'. Apply a 'MonadSample' transformer to
+-- obtain a 'MonadInfer' that can execute probabilistic models.
+module Control.Monad.Bayes.Weighted
+  ( Weighted,
     withWeight,
     runWeighted,
     extractWeight,
@@ -20,18 +18,18 @@
     flatten,
     applyWeight,
     hoist,
-                  ) where
-
-import Control.Monad.Trans
-import Control.Monad.Trans.State
+  )
+where
 
-import Numeric.Log
 import Control.Monad.Bayes.Class
+import Control.Monad.Trans (MonadIO, MonadTrans (..))
+import Control.Monad.Trans.State (StateT (..), mapStateT, modify)
+import Numeric.Log (Log)
 
 -- | Execute the program using the prior distribution, while accumulating likelihood.
 newtype Weighted m a = Weighted (StateT (Log Double) m a)
-    -- StateT is more efficient than WriterT
-    deriving(Functor, Applicative, Monad, MonadIO, MonadTrans, MonadSample)
+  -- StateT is more efficient than WriterT
+  deriving (Functor, Applicative, Monad, MonadIO, MonadTrans, MonadSample)
 
 instance Monad m => MonadCond (Weighted m) where
   score w = Weighted (modify (* w))
@@ -42,27 +40,28 @@
 runWeighted :: (Functor m) => Weighted m a -> m (a, Log Double)
 runWeighted (Weighted m) = runStateT m 1
 
+-- | Compute the sample and discard the weight.
+--
+-- This operation introduces bias.
+prior :: Functor m => Weighted m a -> m a
+prior = fmap fst . runWeighted
+
 -- | Compute the weight and discard the sample.
 extractWeight :: Functor m => Weighted m a -> m (Log Double)
-extractWeight m = snd <$> runWeighted m
+extractWeight = fmap snd . runWeighted
 
 -- | Embed a random variable with explicitly given likelihood.
 --
 -- > runWeighted . withWeight = id
 withWeight :: (Monad m) => m (a, Log Double) -> Weighted m a
 withWeight m = Weighted $ do
-  (x,w) <- lift m
+  (x, w) <- lift m
   modify (* w)
   return x
 
--- | Discard the weight.
--- This operation introduces bias.
-prior :: (Functor m) => Weighted m a -> m a
-prior = fmap fst . runWeighted
-
 -- | Combine weights from two different levels.
 flatten :: Monad m => Weighted (Weighted m) a -> Weighted m a
-flatten m = withWeight $ (\((x,p),q) -> (x, p*q)) <$> runWeighted (runWeighted m)
+flatten m = withWeight $ (\((x, p), q) -> (x, p * q)) <$> runWeighted (runWeighted m)
 
 -- | Use the weight as a factor in the transformed monad.
 applyWeight :: MonadCond m => Weighted m a -> m a
diff --git a/test/Spec.hs b/test/Spec.hs
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -1,18 +1,17 @@
 import Test.Hspec
 import Test.Hspec.QuickCheck
 import Test.QuickCheck
-
-import qualified TestWeighted
 import qualified TestEnumerator
+import qualified TestInference
 import qualified TestPopulation
 import qualified TestSequential
-import qualified TestInference
-
+import qualified TestWeighted
 
 main :: IO ()
 main = hspec $ do
-  describe "Weighted" $
-    it "accumulates likelihood correctly" $ do
+  describe "Weighted"
+    $ it "accumulates likelihood correctly"
+    $ do
       passed <- TestWeighted.passed
       passed `shouldBe` True
   describe "Dist" $ do
@@ -24,7 +23,7 @@
       TestEnumerator.passed4 `shouldBe` True
   describe "Population" $ do
     context "controlling population" $ do
-      it "preserves the population when not expicitly altered" $ do
+      it "preserves the population when not explicitly altered" $ do
         popSize <- TestPopulation.popSize
         popSize `shouldBe` 5
       it "multiplies the number of samples when spawn invoked twice" $ do
@@ -60,9 +59,9 @@
         observations >= 0 && particles >= 1 ==> ioProperty $ do
           checkParticles <- TestInference.checkParticles observations particles
           return $ checkParticles == particles
-  describe "SMC with systematic resampling" $ do
-    prop "number of particles is equal to its second parameter" $
-      \observations particles ->
-        observations >= 0 && particles >= 1 ==> ioProperty $ do
-          checkParticles <- TestInference.checkParticlesSystematic observations particles
-          return $ checkParticles == particles
+  describe "SMC with systematic resampling"
+    $ prop "number of particles is equal to its second parameter"
+    $ \observations particles ->
+      observations >= 0 && particles >= 1 ==> ioProperty $ do
+        checkParticles <- TestInference.checkParticlesSystematic observations particles
+        return $ checkParticles == particles
diff --git a/test/TestEnumerator.hs b/test/TestEnumerator.hs
--- a/test/TestEnumerator.hs
+++ b/test/TestEnumerator.hs
@@ -1,31 +1,30 @@
 module TestEnumerator where
 
+import Control.Monad.Bayes.Class
+import Control.Monad.Bayes.Enumerator
 import Data.AEq
 import qualified Data.Vector as V
-
 import Numeric.Log
-import Control.Monad.Bayes.Enumerator
-import Control.Monad.Bayes.Class
 import Sprinkler
 
 unnorm :: MonadSample m => m Int
-unnorm = categorical $ V.fromList [0.5,0.8]
+unnorm = categorical $ V.fromList [0.5, 0.8]
 
 passed1 :: Bool
 passed1 = (exp . ln) (evidence unnorm) ~== 1
 
 agg :: MonadSample m => m Int
 agg = do
-  x <- uniformD [0,1]
-  y <- uniformD [2,1]
-  return (x+y)
+  x <- uniformD [0, 1]
+  y <- uniformD [2, 1]
+  return (x + y)
 
 passed2 :: Bool
-passed2 = enumerate agg ~== [(1,0.25), (2,0.5), (3,0.25)]
+passed2 = enumerate agg ~== [(1, 0.25), (2, 0.5), (3, 0.25)]
 
 passed3 :: Bool
 passed3 = enumerate Sprinkler.hard ~== enumerate Sprinkler.soft
 
 passed4 :: Bool
 passed4 =
- expectation (^ (2 :: Int)) (fmap (fromIntegral . (+1)) $ categorical $ V.fromList [0.5, 0.5]) ~== 2.5
+  expectation (^ (2 :: Int)) (fmap (fromIntegral . (+ 1)) $ categorical $ V.fromList [0.5, 0.5]) ~== 2.5
diff --git a/test/TestInference.hs b/test/TestInference.hs
--- a/test/TestInference.hs
+++ b/test/TestInference.hs
@@ -1,18 +1,15 @@
-{-# LANGUAGE
-  Rank2Types,
-  TypeFamilies
- #-}
+{-# LANGUAGE Rank2Types #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module TestInference where
 
-import Data.AEq
-import Numeric.Log
-
 import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Enumerator
-import Control.Monad.Bayes.Sampler
-import Control.Monad.Bayes.Population
 import Control.Monad.Bayes.Inference.SMC
+import Control.Monad.Bayes.Population
+import Control.Monad.Bayes.Sampler
+import Data.AEq
+import Numeric.Log
 import Sprinkler
 
 sprinkler :: MonadInfer m => m Bool
@@ -31,5 +28,6 @@
 checkTerminateSMC = sampleIOfixed (runPopulation $ smcMultinomial 2 5 sprinkler)
 
 checkPreserveSMC :: Bool
-checkPreserveSMC = (enumerate . collapse . smcMultinomial 2 2) sprinkler ~==
-                      enumerate sprinkler
+checkPreserveSMC =
+  (enumerate . collapse . smcMultinomial 2 2) sprinkler
+    ~== enumerate sprinkler
diff --git a/test/TestPopulation.hs b/test/TestPopulation.hs
--- a/test/TestPopulation.hs
+++ b/test/TestPopulation.hs
@@ -1,11 +1,10 @@
 module TestPopulation where
 
-import Data.AEq
-
 import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Enumerator
-import Control.Monad.Bayes.Sampler
 import Control.Monad.Bayes.Population as Population
+import Control.Monad.Bayes.Sampler
+import Data.AEq
 import Sprinkler
 
 weightedSampleSize :: MonadSample m => Population m a -> m Int
@@ -26,18 +25,22 @@
 --all_check = (mass (Population.all id (spawn 2 >> sprinkler)) True) ~== 0.09
 
 transCheck1 :: Bool
-transCheck1 = enumerate (collapse sprinkler) ~==
-               sprinklerExact
+transCheck1 =
+  enumerate (collapse sprinkler)
+    ~== sprinklerExact
+
 transCheck2 :: Bool
-transCheck2 = enumerate (collapse (spawn 2 >> sprinkler)) ~==
-               sprinklerExact
+transCheck2 =
+  enumerate (collapse (spawn 2 >> sprinkler))
+    ~== sprinklerExact
 
 resampleCheck :: Int -> Bool
 resampleCheck n =
-  (enumerate . collapse . resampleMultinomial) (spawn n >> sprinkler) ~==
-  sprinklerExact
+  (enumerate . collapse . resampleMultinomial) (spawn n >> sprinkler)
+    ~== sprinklerExact
 
 popAvgCheck :: Bool
-popAvgCheck = expectation f Sprinkler.soft ~== expectation id (popAvg f $ pushEvidence Sprinkler.soft) where
-  f True = 10
-  f False = 4
+popAvgCheck = expectation f Sprinkler.soft ~== expectation id (popAvg f $ pushEvidence Sprinkler.soft)
+  where
+    f True = 10
+    f False = 4
diff --git a/test/TestSequential.hs b/test/TestSequential.hs
--- a/test/TestSequential.hs
+++ b/test/TestSequential.hs
@@ -1,29 +1,29 @@
 module TestSequential where
 
-import Data.AEq
-
 import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Enumerator as Dist
 import Control.Monad.Bayes.Sequential
+import Data.AEq
 import Sprinkler
 
 twoSync :: MonadInfer m => m Int
 twoSync = do
-  x <- uniformD[0,1]
+  x <- uniformD [0, 1]
   factor (fromIntegral x)
-  y <- uniformD[0,1]
+  y <- uniformD [0, 1]
   factor (fromIntegral y)
-  return (x+y)
+  return (x + y)
 
 finishedTwoSync :: MonadInfer m => Int -> m Bool
-finishedTwoSync n = finished (run n twoSync) where
-  run 0 d = d
-  run k d = run (k-1) (advance d)
+finishedTwoSync n = finished (run n twoSync)
+  where
+    run 0 d = d
+    run k d = run (k -1) (advance d)
 
 checkTwoSync :: Int -> Bool
 checkTwoSync 0 = mass (finishedTwoSync 0) False ~== 1
 checkTwoSync 1 = mass (finishedTwoSync 1) False ~== 1
-checkTwoSync 2 = mass (finishedTwoSync 2) True  ~== 1
+checkTwoSync 2 = mass (finishedTwoSync 2) True ~== 1
 checkTwoSync _ = error "Unexpected argument"
 
 sprinkler :: MonadInfer m => m Bool
@@ -39,9 +39,10 @@
 pFinished _ = error "Unexpected argument"
 
 isFinished :: MonadInfer m => Int -> m Bool
-isFinished n = finished (run n sprinkler) where
-  run 0 d = d
-  run k d = run (k-1) (advance d)
+isFinished n = finished (run n sprinkler)
+  where
+    run 0 d = d
+    run k d = run (k -1) (advance d)
 
 checkSync :: Int -> Bool
 checkSync n = mass (isFinished n) True ~== pFinished n
diff --git a/test/TestWeighted.hs b/test/TestWeighted.hs
--- a/test/TestWeighted.hs
+++ b/test/TestWeighted.hs
@@ -1,34 +1,31 @@
-{-# LANGUAGE
-  TypeFamilies
- #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module TestWeighted where
 
-import Data.AEq
-import Control.Monad.State
-import Data.Bifunctor (second)
-import Numeric.Log
-
 import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Sampler
 import Control.Monad.Bayes.Weighted
+import Control.Monad.State
+import Data.AEq
+import Data.Bifunctor (second)
+import Numeric.Log
 
-model :: MonadInfer m => m (Int,Double)
+model :: MonadInfer m => m (Int, Double)
 model = do
-  n <- uniformD [0,1,2]
+  n <- uniformD [0, 1, 2]
   unless (n == 0) (factor 0.5)
   x <- if n == 0 then return 1 else normal 0 1
-  when (n == 2) (factor $ (Exp . log) (x*x))
-  return (n,x)
+  when (n == 2) (factor $ (Exp . log) (x * x))
+  return (n, x)
 
-result :: MonadSample m => m ((Int,Double), Double)
+result :: MonadSample m => m ((Int, Double), Double)
 result = second (exp . ln) <$> runWeighted model
 
 passed :: IO Bool
 passed = fmap check (sampleIOfixed result)
 
-check :: ((Int,Double), Double) -> Bool
-check ((0,1),1) = True
-check ((1,_),y) =  y ~== 0.5
-check ((2,x),y) =  y ~== 0.5 * x * x
+check :: ((Int, Double), Double) -> Bool
+check ((0, 1), 1) = True
+check ((1, _), y) = y ~== 0.5
+check ((2, x), y) = y ~== 0.5 * x * x
 check _ = False