diff --git a/.DS_Store b/.DS_Store
new file mode 100644
Binary files /dev/null and b/.DS_Store differ
diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,7 @@
+Copyright 2021 Sam Staton and others
+
+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
new file mode 100644
--- /dev/null
+++ b/Setup.hs
@@ -0,0 +1,2 @@
+import Distribution.Simple
+main = defaultMain
diff --git a/lazyppl.cabal b/lazyppl.cabal
new file mode 100644
--- /dev/null
+++ b/lazyppl.cabal
@@ -0,0 +1,50 @@
+cabal-version: 1.12
+
+-- This file has been generated from package.yaml by hpack version 0.35.1.
+--
+-- see: https://github.com/sol/hpack
+
+name:           lazyppl
+version:        1.0
+synopsis:       Lazy Probabilistic Programming Library
+description:    LazyPPL is a Haskell library for Bayesian probabilistic programming. It supports lazy use of probability, which is useful for specifying non-parametric models, and we provide new Metropolis-Hastings algorithms to allow this. For illustrations, see <https://lazyppl-team.github.io/>. The Gaussian Process module uses hmatrix, which requires a separate lapack installation (or -dev or -devel).
+category:       Statistics
+stability:      experimental
+homepage:       https://lazyppl-team.github.io/
+bug-reports:    https://github.com/lazyppl-team/lazyppl/issues
+author:         LazyPPL team
+maintainer:     sam.staton@cs.ox.ac.uk
+copyright:      2021-2024 LazyPPL team
+license:        MIT
+license-file:   LICENSE
+build-type:     Simple
+
+source-repository head
+  type: git
+  location: https://github.com/lazyppl-team/lazyppl
+
+library
+  exposed-modules:
+      LazyPPL
+      LazyPPL.Distributions
+      LazyPPL.Distributions.Counter
+      LazyPPL.Distributions.DirichletP
+      LazyPPL.Distributions.GP
+      LazyPPL.Distributions.IBP
+      LazyPPL.Distributions.Memoization
+  other-modules:
+      Paths_lazyppl
+  hs-source-dirs:
+      src
+  build-depends:
+      base >=4.7 && <5
+    , containers >=0.6.5.1 && <0.7
+    , hmatrix >=0.20.2 && <0.21
+    , log-domain >=0.13.2 && <0.14
+    , math-functions >=0.3.4.2 && <0.3.5
+    , monad-extras >=0.6.0 && <0.7
+    , mtl >=2.2.2 && <2.4
+    , random >=1.2.0 && <1.3
+    , transformers >=0.5.6.2 && <0.7
+  default-language: Haskell2010
+
diff --git a/src/.DS_Store b/src/.DS_Store
new file mode 100644
Binary files /dev/null and b/src/.DS_Store differ
diff --git a/src/LazyPPL.hs b/src/LazyPPL.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL.hs
@@ -0,0 +1,333 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
+
+{- | LazyPPL is a library for Bayesian probabilistic programming. It supports lazy use of probability, and we provide new Metropolis-Hastings simulation algorithms to allow this. Laziness appears to be a good paradigm for non-parametric statistics. 
+
+Reference paper: [Affine Monads and Lazy Structures for Bayesian Programming](https://arxiv.org/abs/2212.07250). POPL 2023.
+
+Illustrations: [https://lazyppl-team.github.io](https://lazyppl-team.github.io).
+
+
+LazyPPL is inspired by recent ideas in synthetic probability theory and synthetic measure theory, such as [quasi-Borel spaces](https://ncatlab.org/nlab/show/quasi-Borel+space) and [Markov categories](https://ncatlab.org/nlab/show/Markov+category). LazyPPL is inspired by many other languages, including [Church](http://v1.probmods.org), [Anglican](https://probprog.github.io/anglican/), and [Monad-Bayes](https://hackage.haskell.org/package/monad-bayes). Monad-Bayes now includes a LazyPPL-inspired simulation algorithm.
+
+This module defines
+
+    1. Two monads: `Prob` (for probability measures) and `Meas` (for unnormalized measures), with interface `uniform`, `sample`, `score`. 
+
+    2. Monte Carlo inference methods produce samples from an unnormalized measure. We provide three inference methods: 
+
+        a. 'mh' (Metropolis-Hastings algorithm based on lazily mutating parts of the tree at random).
+
+        b. 'mhirreducible', which randomly restarts for a properly irreducible Metropolis-Hastings kernel.
+
+        c. 'wis' (simple reference weighted importance sampling)
+
+        See also the SingleSite module for a separate single-site Metropolis-Hastings algorithm via GHC.Exts.Heap and System.IO.Unsafe.
+
+    3. Various useful helpful functions.
+
+    A typical usage would be
+
+@   
+    import LazyPPL (Prob, Meas, uniform, sample, score, mh, every)
+@
+
+    Most of the structure here will not be needed in typical models. We expose more of the structure for more experimental uses. 
+
+The `Distributions` module provides many useful distributions, and further non-parametric distributions are in `Distributions.DirichletP`, `Distributions.GP`, `Distr.IBP`, and `Distr.Memoization`. 
+
+
+-}
+
+module LazyPPL
+    ( -- * Rose tree type
+      --
+      -- | Our source of randomness will be an infinitely wide and deep lazy [rose tree](https://en.wikipedia.org/wiki/Rose_tree), regarded as initialized with uniform [0,1] choices for each label.
+      Tree(Tree),
+      -- * Monads
+      Prob(Prob), Meas(Meas),
+      -- * Basic interface
+      --
+      -- | There are three building blocks for measures: `uniform` for probability measures; `sample` and `score` for unnormalized measures. Combined with the monad structure, these give all s-finite measures.
+      uniform, sample, score,
+      -- * Monte Carlo simulation
+      --
+      -- | The `Meas` type describes unnormalized measures. Monte Carlo simulation allows us to sample from an unnormalized measure. Our main Monte Carlo simulator is `mh`. 
+      mh, mhirreducible, weightedsamples, wis, 
+      -- * Useful functions
+      every, randomTree, runProb) where
+
+import Control.Monad.Trans.Writer
+import Control.Monad.Trans.Class
+import Data.Monoid
+import System.Random hiding (uniform)
+import Control.Monad
+import Control.Monad.Extra
+import Control.Monad.State.Lazy (State, state , put, get, runState)
+import Numeric.Log
+
+
+{- | A `Tree` here is a lazy, infinitely wide and infinitely deep rose tree, labelled by Doubles.
+-}
+data Tree = Tree Double [Tree]
+-- Often people would just use a list or stream instead of a tree.
+-- But a tree allows us to be lazy about how far we are going all the time.
+
+{- | A probability distribution over a is 
+a function @ Tree -> a @.
+
+We can think of this as the law of a random variable, indexed by the source of randomness, which is `Tree`. 
+
+According to the monad implementation, a program uses up bits of the tree as it runs. The tree being infinitely wide and deep allows for lazy computation.-}
+newtype Prob a = Prob (Tree -> a)
+
+-- | Split tree splits a tree in two (bijectively)
+splitTree :: Tree -> (Tree , Tree)
+splitTree (Tree r (t : ts)) = (t , Tree r ts)
+
+
+-- | Sequencing is done by splitting the tree
+-- and using different bits for different computations.
+-- 
+-- This monad structure is strongly inspired by the probability monad of [quasi-Borel space](https://ncatlab.org/nlab/show/quasi-Borel+space#probability_distributions). 
+instance Monad Prob where
+  return a = Prob $ const a
+  (Prob m) >>= f = Prob $ \g ->
+    let (g1, g2) = splitTree g
+        (Prob m') = f (m g1)
+    in m' g2
+instance Functor Prob where fmap = liftM
+instance Applicative Prob where {pure = return ; (<*>) = ap}
+
+{- | An unnormalized measure is represented by a probability distribution over pairs of a weight and a result. -}
+newtype Meas a = Meas (WriterT (Product (Log Double)) Prob a)
+  deriving(Functor, Applicative, Monad)
+
+{- | A uniform sample is a building block for probability distributions.
+
+This is implemented by getting the label at the head of the tree and discarding the rest.-}
+uniform :: Prob Double
+uniform = Prob $ \(Tree r _) -> r
+
+-- | Regard a probability measure as an unnormalized measure.
+sample :: Prob a -> Meas a
+sample p = Meas $ lift p
+
+{- | A one point measure with a given score (or weight, or mass, or likelihood), which should be a positive real number.
+
+A score of 0 describes impossibility. To avoid numeric issues, we encode it as @ exp(-300) @ instead.-}
+score :: Double -> Meas ()
+score r = Meas $ tell $ Product $ (Exp . log) (if r==0 then exp(-300) else r)
+
+scoreLog :: Log Double -> Meas ()
+scoreLog r = Meas $ tell $ Product r
+
+scoreProductLog :: Product (Log Double) -> Meas ()
+scoreProductLog r = Meas $ tell r
+
+
+{- | Generate a tree with uniform random labels.
+
+    This uses 'split' to split a random seed. -}
+randomTree :: RandomGen g => g -> Tree
+randomTree g = let (a,g') = random g in Tree a (randomTrees g')
+randomTrees :: RandomGen g => g -> [Tree]
+randomTrees g = let (g1,g2) = split g in randomTree g1 : randomTrees g2
+
+{- | 'runProb' runs a probability deterministically, given a source of randomness. -}
+runProb :: Prob a -> Tree -> a
+runProb (Prob a) = a
+
+{- | Runs an unnormalized measure and gets out a stream of (result,weight) pairs.
+
+These are not samples from the renormalized distribution, just plain (result,weight) pairs. This is useful when the distribution is known to be normalized already. -}
+weightedsamples :: forall a. Meas a -> IO [(a,Log Double)]
+weightedsamples (Meas m) =
+  do
+    let helper :: Prob [(a, Product (Log Double))]
+        helper = do
+          (x, w) <- runWriterT m
+          rest <- helper
+          return $ (x, w) : rest
+    newStdGen
+    g <- getStdGen
+    let rs = randomTree g
+    let xws = runProb helper rs
+    return $ map (\(x,w) -> (x, getProduct w)) xws
+
+{- | Weighted importance sampling first draws n weighted samples,
+    and then samples a stream of results from that, regarded as an empirical distribution. Sometimes called "likelihood weighted importance sampling". 
+
+This is a reference implementation. It will not usually be very efficient at all, but may be useful for debugging. 
+ -}
+wis :: Int -- ^ @n@, the number of samples to base on
+    -> Meas a -- ^ @m@, the measure to normalize
+    -> IO [a] -- ^ Returns a stream of samples
+wis n m = do
+  xws <- weightedsamples m
+  let xws' = take n $ accumulate xws 0
+  let max = snd $ last xws'
+  newStdGen
+  g <- getStdGen
+  let rs = (randoms g :: [Double])
+  return $ map (\r -> fst $ head $ filter (\(x, w) -> w >= Exp (log r) * max) xws') rs
+  where accumulate ((x, w) : xws) a = (x, w + a) : (x, w + a) : accumulate xws (w + a)
+        accumulate [] a = []
+
+{- | Produce a stream of samples, using Metropolis Hastings simulation.
+
+   The weights are also returned. Often the weights can be discarded, but sometimes we may search for a sample of maximum score.
+
+   The algorithm works as follows. 
+
+   At each step, we randomly change some sites (nodes in the tree). 
+   We then accept or reject these proposed changes, using a probability that is determined by the weight of the measure at the new tree. 
+   If rejected, we repeat the previous sample. 
+
+    This kernel is related to the one introduced by [Wingate, Stuhlmuller, Goodman, AISTATS 2011](http://proceedings.mlr.press/v15/wingate11a/wingate11a.pdf), but it is different in that it works when the number of sites is unknown. Moreover, since a site is a path through the tree, the address is more informative than a number, which avoids some addressing issues. 
+
+    When 1/@p@ is roughly the number of used sites, then this will be a bit like "single-site lightweight" MH.
+    If @p@ = 1 then this is "multi-site lightweight" MH.
+
+    __Tip:__ if @m :: Prob a@ then use @map fst <$> (mh 1 $ sample m)@ to get a stream of samples from a probability distribution without conditioning. 
+--}
+{-- The algorithm is as follows:
+
+    Top level: produces a stream of samples.
+    * Split the random number generator in two
+    * One part is used as the first seed for the simulation,
+    * and one part is used for the randomness in the MH algorithm.
+
+    Then, run 'step' over and over to get a stream of '(tree, result, weight)'s.
+    The stream of seeds is used to produce a stream of result/weight pairs.
+
+    NB There are three kinds of randomness in the step function.
+    1. The start tree 't', which is the source of randomness for simulating the
+    program m to start with. This is sort-of the point in the "state space".
+    2. The randomness needed to propose a new tree ('g1')
+    3. The randomness needed to decide whether to accept or reject that ('g2')
+    The tree t is an argument and result,
+    but we use a state monad ('get'/'put') to deal with the other randomness '(g, g1, g2)'
+
+    Steps of the 'step' function:
+    1. Randomly change some sites, 
+    2. Rerun the model with the new tree, to get a new weight 'w''.
+    3. Compute the MH acceptance ratio. This is the probability of either returning the new seed or the old one.
+    4. Accept or reject the new sample based on the MH ratio.  
+--}
+mh :: forall a.
+   Double -- ^ The chance @p@ of changing any site
+   -> Meas a -- ^ The unnormalized measure to sample from
+   -> IO [(a,Product (Log Double))] -- ^ Returns a stream of (result,weight) pairs
+mh p (Meas m) = do
+    -- Top level: produce a stream of samples.
+    -- Split the random number generator in two
+    -- One part is used as the first seed for the simulation,
+    -- and one part is used for the randomness in the MH algorithm.
+    newStdGen
+    g <- getStdGen
+    let (g1,g2) = split g
+    let t = randomTree g1
+    let (x, w) = runProb (runWriterT m) t
+    -- Now run step over and over to get a stream of (tree,result,weight)s.
+    let (samples,_) = runState (iterateM step (t,x,w)) g2
+    -- The stream of seeds is used to produce a stream of result/weight pairs.
+    return $ map (\(_,x,w) -> (x,w)) samples
+    {- NB There are three kinds of randomness in the step function.
+    1. The start tree 't', which is the source of randomness for simulating the
+    program m to start with. This is sort-of the point in the "state space".
+    2. The randomness needed to propose a new tree ('g1')
+    3. The randomness needed to decide whether to accept or reject that ('g2')
+    The tree t is an argument and result,
+    but we use a state monad ('get'/'put') to deal with the other randomness '(g,g1,g2)' -}
+    where step :: RandomGen g => (Tree,a,Product (Log Double)) -> State g (Tree,a,Product (Log Double))
+          step (t, x, w) = do
+            -- Randomly change some sites
+            g <- get
+            let (g1, g2) = split g
+            let t' = mutateTree p g1 t
+            -- Rerun the model with the new tree, to get a new
+            -- weight w'.
+            let (x', w') = runProb (runWriterT m) t'
+            -- MH acceptance ratio. This is the probability of either
+            -- returning the new seed or the old one.
+            let ratio = getProduct w' / getProduct w
+            let (r, g2') = random g2
+            put g2'
+            if r < min 1 (exp $ ln ratio) -- (trace ("-- Ratio: " ++ show ratio) ratio))
+              then return (t', x', w') -- trace ("---- Weight: " ++ show w') w')
+              else return (t, x, w) -- trace ("---- Weight: " ++ show w) w)
+
+
+-- | Replace the labels of a tree randomly, with probability p
+mutateTree :: forall g. RandomGen g => Double -> g -> Tree -> Tree
+mutateTree p g (Tree a ts) =
+  let (a',g') = (random g :: (Double,g)) in
+  let (a'',g'') = random g' in
+  Tree (if a'<p then a'' else a) (mutateTrees p g'' ts)
+mutateTrees :: RandomGen g => Double -> g -> [Tree] -> [Tree]
+mutateTrees p g (t:ts) = let (g1,g2) = split g in mutateTree p g1 t : mutateTrees p g2 ts
+
+
+
+{- | Irreducible form of 'mh'. Takes @p@ like 'mh', but also @q@, which is the chance of proposing an all-sites change. Irreducibility means that, asymptotically, the sequence of samples will converge in distribution to the renormalized version of @m@. 
+
+The kernel in `mh` is not formally irreducible in the usual sense, although it is an open question whether this is a problem for asymptotic convergence in any definable model. In any case, convergence is only asymptotic, and so it can be helpful to use `mhirreducible` is that in some situations.
+
+Roughly this avoids `mh` getting stuck in one particular mode, although it is a rather brutal method.
+
+ -}
+
+mhirreducible :: forall a.
+              Double -- ^ The chance @p@ of changing any given site
+              -> Double -- ^ The chance @q@ of doing an all-sites change
+              -> Meas a -- ^ The unnormalized measure @m@ to sample from
+              -> IO [(a,Product (Log Double))] -- ^ Returns a stream of (result,weight) pairs
+mhirreducible p q (Meas m) = do
+    -- Top level: produce a stream of samples.
+    -- Split the random number generator in two
+    -- One part is used as the first seed for the simulation,
+    -- and one part is used for the randomness in the MH algorithm.
+    newStdGen
+    g <- getStdGen
+    let (g1,g2) = split g
+    let t = randomTree g1
+    let (x, w) = runProb (runWriterT m) t
+    -- Now run step over and over to get a stream of (tree,result,weight)s.
+    let (samples,_) = runState (iterateM step (t,x,w)) g2
+    -- The stream of seeds is used to produce a stream of result/weight pairs.
+    return $ map (\(t,x,w) -> (x,w)) samples
+    {- NB There are three kinds of randomness in the step function.
+    1. The start tree 't', which is the source of randomness for simulating the
+    program m to start with. This is sort-of the point in the "state space".
+    2. The randomness needed to propose a new tree ('g1')
+    3. The randomness needed to decide whether to accept or reject that ('g2')
+    The tree t is an argument and result,
+    but we use a state monad ('get'/'put') to deal with the other randomness '(g,g1,g2)' -}
+    where step :: RandomGen g => (Tree,a,Product (Log Double)) -> State g (Tree,a,Product (Log Double))
+          step (t,x,w) = do
+            -- Randomly change some sites
+            g <- get
+            let (g1,g2) = split g
+            -- Decide whether to resample all sites (r<q) or just some of them
+            let (r,g1') = random g1
+            let t' = if r<q then randomTree g1' else mutateTree p g1' t
+            -- Rerun the model with the new tree, to get a new
+            -- weight w'.
+            let (x', w') = runProb (runWriterT m) t'
+            -- MH acceptance ratio. This is the probability of either
+            -- returning the new seed or the old one.
+            let ratio = getProduct w' / getProduct w
+            let (r,g2') = random g2
+            put g2'
+            if r < min 1 (exp $ ln ratio) then return (t',x',w') else return (t,x,w)
+
+
+
+{- | Useful function which thins out a stream of results, as is common in Markov Chain Monte Carlo simulation.
+
+@every n xs@ returns only the elements at indices that are multiples of n.--}
+every :: Int -> [a] -> [a]
+every n xs = case drop (n -1) xs of
+  (y : ys) -> y : every n ys
+  [] -> []
+
diff --git a/src/LazyPPL/.DS_Store b/src/LazyPPL/.DS_Store
new file mode 100644
Binary files /dev/null and b/src/LazyPPL/.DS_Store differ
diff --git a/src/LazyPPL/Distributions.hs b/src/LazyPPL/Distributions.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions.hs
@@ -0,0 +1,132 @@
+{- | Handy distributions for the `LazyPPL` library, based on the `uniform` distribution. Mostly defined using the `Statistics.Distribution` module and family. 
+
+Sometimes both a distribution (type @Prob a@) and pdf (type @a -> Double@) are given. Distributions are useful for sampling, densities are used for scoring. 
+
+For more distributions, see the Statistics.Distribution in the statistics package. -}
+
+
+module LazyPPL.Distributions (
+       -- * Continuous distributions
+       normal,normalPdf,exponential,expPdf,gamma, beta, dirichlet, uniformbounded,
+       -- * Discrete distributions
+       bernoulli, uniformdiscrete, categorical, poisson, poissonPdf,
+       -- * Streams
+       iid)
+       where
+
+import LazyPPL (Prob,uniform)
+import Data.List (findIndex)
+import Numeric.SpecFunctions
+import Numeric.MathFunctions.Constants
+
+{-|
+  [Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution)
+-}
+normal :: Double -- ^ mu, mean
+       -> Double -- ^ sigma, standard deviation
+       -> Prob Double
+normal m s = do 
+  x <- uniform
+  return $ (- invErfc (2 * x)) * (m_sqrt_2 * s) + m
+
+normalPdf :: Double -> Double -> Double -> Double
+normalPdf m s x = exp ((-(x - m) * (x -m) / (2 * s * s)) - log (m_sqrt_2_pi * s))
+
+
+{-|
+  [Exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution)
+-}
+exponential :: Double -- ^ lambda, rate
+            -> Prob Double
+exponential rate = do 
+  x <- uniform
+  return $ - (log x / rate)
+
+expPdf :: Double -> Double -> Double
+expPdf rate x = exp (-rate*x) * rate
+
+{-|
+  [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution)
+-}
+gamma :: Double -- ^ k, shape
+      -> Double -- ^ theta, scale
+      -> Prob Double
+gamma a b = do
+  x <- uniform
+  return $ b * invIncompleteGamma a x
+
+{-|
+  [Beta distribution](https://en.wikipedia.org/wiki/Beta_distribution)
+-}
+beta :: Double -- ^ alpha
+     -> Double -- ^ beta
+     -> Prob Double
+beta a b = do
+  x <- uniform
+  return $ invIncompleteBeta a b x
+
+{-|
+  [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution)
+-}
+poisson :: Double -- ^ lambda, rate
+        -> Prob Integer
+poisson lambda = do
+  x <- uniform
+  let cmf = map (\x -> 1 - incompleteGamma (fromIntegral (x + 1)) lambda) [0,1..]
+  let (Just n) = findIndex (> x) cmf
+  return $ fromIntegral n
+
+poissonPdf :: Double -> Integer -> Double
+poissonPdf rate n = let result = exp(-rate) * rate ^^ (fromIntegral n) / (factorial (fromIntegral n)) in 
+  if (isInfinite result) || (isNaN result) then exp (-rate + (fromIntegral n) * log rate - logGamma (fromIntegral (n+1))) else result
+
+
+{-|
+  [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
+-}
+dirichlet :: [Double] -- ^ vector of alphas; length is dimension
+          -> Prob[Double]
+dirichlet as = do
+  xs <- mapM (\a -> gamma a 1) as
+  let s = Prelude.sum xs
+  let ys = map (/ s) xs
+  return ys
+
+-- | [Continuous uniform distribution on a bounded interval](https://en.wikipedia.org/wiki/Continuous_uniform_distribution)
+uniformbounded :: Double -- ^ lower
+               -> Double -- ^ upper
+               -> Prob Double
+uniformbounded lower upper = do
+  x <- uniform
+  return $ (upper - lower) * x + lower
+
+-- | [Bernoulli distribution](https://en.wikipedia.org/wiki/Bernoulli_distribution)
+bernoulli :: Double -- ^ bias
+          -> Prob Bool
+bernoulli r = do
+  x <- uniform
+  return $ x < r
+
+{-|
+  [Discrete uniform distribution](https://en.wikipedia.org/wiki/Discrete_uniform_distribution) on [0, ..., n-1]
+-}
+uniformdiscrete :: Int -- ^ n
+                -> Prob Int
+uniformdiscrete n =
+  do
+    let upper = fromIntegral n
+    r <- uniformbounded 0 upper
+    return $ floor r
+
+{-| [Categorical distribution](https://www.google.com/search?client=safari&rls=en&q=categorical+distribution&ie=UTF-8&oe=UTF-8): Takes a list of k numbers that sum to 1, 
+    and returns a random number between 0 and (k-1), weighted accordingly -}
+categorical :: [Double] -> Prob Int
+categorical xs = do 
+  r <- uniform
+  case findIndex (>r) $ tail $ scanl (+) 0 xs of
+    Just i -> return i
+    Nothing -> error "categorical: probabilities do not sum to 1"
+
+{-| Returns an infinite stream of samples from the given distribution. --}
+iid :: Prob a -> Prob [a]
+iid p = do r <- p; rs <- iid p; return $ r : rs
diff --git a/src/LazyPPL/Distributions/.DS_Store b/src/LazyPPL/Distributions/.DS_Store
new file mode 100644
Binary files /dev/null and b/src/LazyPPL/Distributions/.DS_Store differ
diff --git a/src/LazyPPL/Distributions/Counter.hs b/src/LazyPPL/Distributions/Counter.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions/Counter.hs
@@ -0,0 +1,33 @@
+{-# LANGUAGE BangPatterns #-}
+module LazyPPL.Distributions.Counter (Counter,newCounter,readAndIncrement) where
+
+import LazyPPL
+import Data.IORef
+import System.IO.Unsafe
+
+{--
+Some "unsafe" functions for a hidden counter. 
+This is useful for implementing some nonparametric models 
+like the indian buffet process. 
+The function is deterministic but we perform some redundant sampling
+so that the function is treated as impure by the compiler and 
+thus re-evalued every time.  
+
+The implementation of the counter (IORef Int) is encapsulated.
+--}
+
+data Counter = C (IORef Int)
+
+newCounter :: Prob Counter
+newCounter = do r <- uniform
+                return $ C $ unsafePerformIO $ newIORef (round (r-r))
+
+readAndIncrement :: Counter -> Prob Int 
+readAndIncrement (C ref) = do
+    r <- uniform
+    return $ unsafePerformIO $ do 
+        !i <- readIORef ref 
+        () <- writeIORef ref (i + 1 + round (r - r))
+        return i 
+
+
diff --git a/src/LazyPPL/Distributions/DirichletP.hs b/src/LazyPPL/Distributions/DirichletP.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions/DirichletP.hs
@@ -0,0 +1,67 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+{-| Abstract types for the [Dirichlet Process](https://en.wikipedia.org/wiki/Dirichlet_process) viewed through the interface of the [Chinese Restaurant Process](https://en.wikipedia.org/wiki/Chinese_restaurant_process).
+
+Ideas following [S. Staton, H. Yang, N. L. Ackerman, C. Freer, D. Roy. Exchangeable random process and data abstraction. Workshop on probabilistic programming semantics (PPS 2017).](https://www.cs.ox.ac.uk/people/hongseok.yang/paper/pps17a.pdf)
+
+Our implementation here uses stick breaking, with a lazily broken stick. Other urn-based implementations are possible with hidden state, and they should be observationally equivalent.
+
+For illustrations, see [non-parametric clustering](https://lazyppl-team.github.io/ClusteringDemo.html) and [relational inference](https://lazyppl-team.github.io/IrmDemo.html).
+
+-}
+
+module LazyPPL.Distributions.DirichletP (
+{- * Chinese Restaurant Process interface -}
+{- | For clustering, we regard each data point as a "customer" in a "restaurant", and they are in the same cluster if they sit at the same `Table`. 
+-}
+Restaurant, Table, newRestaurant, newCustomer, 
+-- * Random distribution interface
+dp) where
+
+import Data.List
+import Data.Maybe
+import LazyPPL
+import LazyPPL.Distributions
+import LazyPPL.Distributions.Memoization (MonadMemo)
+
+
+-- | Abstract type of restaurants
+newtype Restaurant = R [Double]
+
+-- | Abstract type of tables. This supports `Eq` so that we can ask whether customers are at the same table (i.e. whether points are in the same cluster). 
+newtype Table = T Int deriving (Eq, Show, MonadMemo Prob)
+
+{-| A customer enters the restaurant and is assigned a table. -}
+newCustomer :: Restaurant -> Prob Table
+newCustomer (R restaurant) =
+  do
+    r <- uniform
+    return $ T $ fromJust $ findIndex (> r) (scanl1 (+) restaurant)
+
+{-| Create a new restaurant with concentration parameter alpha. -}
+newRestaurant :: Double -- ^ Concentration parameter, alpha
+              -> Prob Restaurant
+newRestaurant alpha = do
+  sticks <- stickBreaking alpha 0
+  return $ R sticks
+
+{- | Stick breaking breaks the unit interval into an
+    infinite number of parts (lazily) --}
+stickBreaking :: Double -> Double -> Prob [Double]
+stickBreaking alpha lower =
+  do
+    r <- beta 1 alpha
+    let v = r * (1 - lower)
+    vs <- stickBreaking alpha (lower + v)
+    return (v : vs)
+
+{-| [Dirichlet Process](https://en.wikipedia.org/wiki/Dirichlet_process) as a random distribution. -}
+dp :: Double -- ^ Concentration parameter, alpha
+   -> Prob a -- ^ Base distribution
+   -> Prob (Prob a)
+dp alpha p = do
+  xs <- iid p
+  vs <- stickBreaking alpha 0
+  return $ do
+    r <- uniform
+    return $ xs !! fromJust (findIndex (> r) (scanl1 (+) vs))
diff --git a/src/LazyPPL/Distributions/GP.hs b/src/LazyPPL/Distributions/GP.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions/GP.hs
@@ -0,0 +1,95 @@
+{- | Gaussian processes in the `LazyPPL` library.
+
+Gaussian processes are random functions. We provide the Wiener process (`wiener`) and a generic Gaussian process (`gp`). 
+
+Although the implementation uses hidden state, it is still safe, i.e. statistically commutative and discardable. 
+
+For illustrations, see the [Gaussian process regression](https://lazyppl-team.github.io/GaussianProcessDemo.html) 
+or [Wiener process regression](https://lazyppl-team.github.io/WienerDemo.html). The latter also illustrates composing the Wiener process function to get jump diffusion.
+-}
+
+
+module LazyPPL.Distributions.GP (wiener, gp) where
+
+import LazyPPL (Prob,Tree(Tree),Prob(Prob),runProb)
+import LazyPPL.Distributions (normal,iid)
+import Numeric.LinearAlgebra hiding (step,size,find)
+import Numeric.LinearAlgebra.Data hiding (step,size,find)
+import Data.List (find)
+import Data.Map (empty,lookup,insert,size,keys,elems,Map)
+import Data.IORef
+import System.IO.Unsafe
+import Prelude hiding ((<>))
+
+{-| [Wiener process](https://en.wikipedia.org/wiki/Wiener_process) (Brownian motion).
+
+This is a random function. 
+
+The Wiener process is implemented by using a Brownian bridge and a hidden memo table. -}
+wiener :: Prob (Double -> Double)
+wiener = Prob $ \(Tree r gs) ->
+                   unsafePerformIO $ do
+                                 ref <- newIORef Data.Map.empty
+                                 modifyIORef' ref (Data.Map.insert 0 0)
+                                 return $ \x -> unsafePerformIO $ do
+                                        table <- readIORef ref
+                                        case Data.Map.lookup x table of
+                                             Just y -> do {return y}
+                                             Nothing -> do let lower = do {l <- findMaxLower x (keys table) ;
+                                                                           v <- Data.Map.lookup l table ; return (l,v) }
+                                                           let upper = do {u <- find (> x) (keys table) ;
+                                                                           v <- Data.Map.lookup u table ; return (u,v) }
+                                                           let m = bridge lower x upper
+                                                           let y = runProb m (gs !! (1 + size table))
+                                                           modifyIORef' ref (Data.Map.insert x y)
+                                                           return y
+
+bridge :: Maybe (Double,Double) -> Double -> Maybe (Double,Double) -> Prob Double
+-- not needed since the table is always initialized with (0, 0)
+-- bridge Nothing y Nothing = if y==0 then return 0 else normal 0 (sqrt y) 
+bridge (Just (x,x')) y Nothing = normal x' (sqrt (y-x))
+bridge Nothing y (Just (z,z')) = normal z' (sqrt (z-y))
+bridge (Just (x,x')) y (Just (z,z')) = normal (x' + ((y-x)*(z'-x')/(z-x))) (sqrt ((z-y)*(y-x)/(z-x)))
+
+findMaxLower :: Double -> [Double] -> Maybe Double 
+findMaxLower d [] = Nothing
+findMaxLower d (x:xs) = let y = findMaxLower d xs in
+                       case y of 
+                           Nothing -> if x < d then Just x else Nothing 
+                           Just m -> do 
+                                          if x > m && x < d then Just x else Just m 
+
+
+
+
+{-| [Gaussian process](https://en.wikipedia.org/wiki/Gaussian_process) with given covariance function.
+
+This is a random function.
+
+The function is defined by using linear algebra and a hidden memo table.
+Although it uses hidden state, it is still safe, i.e. statistically commutative and discardable. --}
+gp :: (Double -> Double -> Double) -- ^ Covariance function
+      -> Prob (Double -> Double) -- ^ Returns a random function
+gp cov = do ns <- iid $ normal 0 1
+            return $ unsafePerformIO $ do
+                                 ref <- newIORef Data.Map.empty
+                                 modifyIORef' ref (Data.Map.insert 0 0)
+                                 return $ \x -> unsafePerformIO $ do
+                                        table <- readIORef ref
+                                        case Data.Map.lookup x table of
+                                             Just y -> do {return y}
+                                             Nothing -> do let y = step cov table x $ ns !! (1 + size table)
+                                                           modifyIORef' ref (Data.Map.insert x y)
+                                                           return y
+
+step :: (Double -> Double -> Double) -> Data.Map.Map Double Double -> Double -> Double -> Double
+step cov table x seed =
+       let sig11 = cov x x 
+           sig12 = matrix (size table) [cov x b | b <-keys table] 
+           sig21 = matrix 1 [cov a x | a <-keys table] 
+           sig22 = matrix (size table) [cov a b | a <- keys table , b <-keys table] 
+           regCoeff = sig12 <> (pinvTol 1E8 sig22)
+           mu = (regCoeff <> (matrix 1 $ elems table)) `atIndex` (0,0)
+           var = sig11 - ((regCoeff <> sig21) `atIndex` (0,0)) in
+       mu + seed * (sqrt $ if var > -0.01 then (abs var) else error (show var))
+
diff --git a/src/LazyPPL/Distributions/IBP.hs b/src/LazyPPL/Distributions/IBP.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions/IBP.hs
@@ -0,0 +1,89 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+{-|
+An implementation of the Indian buffet process by [Griffiths and Ghahramani](https://papers.nips.cc/paper_files/paper/2005/file/2ef35a8b78b572a47f56846acbeef5d3-Paper.pdf).
+
+We are using abstract types to hide the implementation details, inspired by [Exchangeable Random Processes and Data Abstraction](https://www.cs.ox.ac.uk/people/hongseok.yang/paper/pps17a.pdf). 
+
+Illustration: [Feature extraction example](https://lazyppl-team.github.io/AdditiveClusteringDemo.html). 
+-} 
+
+module LazyPPL.Distributions.IBP where
+
+import LazyPPL
+import LazyPPL.Distributions
+import LazyPPL.Distributions.Counter
+import LazyPPL.Distributions.Memoization
+
+import Data.List
+
+
+ 
+-- Some abstract types 
+newtype Restaurant = R ([[Bool]], Counter)  
+newtype Dish = D Int  deriving (Eq,Ord,Show,MonadMemo Prob)
+
+newCustomer :: Restaurant -> Prob [Dish]
+newCustomer (R (matrix, ref)) = do 
+    i <- readAndIncrement ref 
+    return [ D k | k <- [0..(length (matrix!!i) - 1)], matrix!!i!!k ]
+
+            
+newRestaurant :: Double -> Prob Restaurant 
+newRestaurant alpha = do
+        r <- uniform 
+        ref <- newCounter
+        matrix <- ibp alpha  
+        return $ R (matrix, ref) 
+
+
+matrix :: Double -> Int -> [Int] -> Prob [[Bool]] 
+matrix alpha index features = 
+     do 
+        let i = fromIntegral index
+        existingDishes <- mapM (\m -> bernoulli ((fromIntegral m) / i)) features 
+        let newFeatures = zipWith (\a -> \b -> if b then a + 1 else a) features existingDishes 
+        nNewDishes     <- fmap fromIntegral $ poisson (alpha / i) 
+        let fixZero = if features == [] && nNewDishes == 0 then 1 else nNewDishes  
+        let newRow = existingDishes ++ (take fixZero $ repeat True) 
+        rest           <- matrix alpha (index + 1) (newFeatures ++ (take fixZero $ repeat 1)) 
+        return $ newRow : rest  
+
+-- the distribution on matrices 
+ibp :: Double -> Prob [[Bool]]  
+ibp alpha = matrix alpha 1 [] 
+
+
+
+{--
+Another possible implementation of the indian buffet process 
+which uses a truncated stickbreaking construction. 
+It is only an approximation to the true IBP, but doesn't need IO.   
+
+See also 
+Stick-breaking Construction for the Indian Buffet Process
+Teh, Gorur, Ghahramani. AISTATS 2007.
+
+A stochastic programming perspective on nonparametric Bayes
+Daniel M. Roy, Vikash Mansinghka, Noah Goodman, and Joshua Tenenbaum
+ICML Workshop on Nonparametric Bayesian, 2008. 
+--}
+data RestaurantS = RS [Double] 
+
+data DishS = DS Int deriving (Eq,Ord,Show)
+
+newCustomerS :: RestaurantS -> Prob [DishS]
+newCustomerS (RS rs) = 
+    do fs <- mapM bernoulli rs
+       return $ map DS $ findIndices id fs
+
+newRestaurantS :: Double -> Prob RestaurantS 
+newRestaurantS a = fmap RS $ stickScale 1
+  where stickScale p = do r' <- beta a 1
+                          let r = p * r'
+                          -- Truncate when the probabilities are getting small
+                          rs <- if r < 0.01 then return [] else stickScale r
+                          return $ r : rs
+
+
+
diff --git a/src/LazyPPL/Distributions/Memoization.hs b/src/LazyPPL/Distributions/Memoization.hs
new file mode 100644
--- /dev/null
+++ b/src/LazyPPL/Distributions/Memoization.hs
@@ -0,0 +1,106 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE DefaultSignatures #-}
+
+{- | Stochastic memoization for the `LazyPPL` library.
+Stochastic memoization is a useful primitive for programming non-parametric models. 
+It has type @(a -> Prob b) -> Prob (a -> b)@, and can be thought of as simultaneously sampling results for all possible arguments, in a lazy way. 
+
+When @a@ is enumerable, this amounts to converting a stream of probabilities to a random stream, which we can do by sampling each probability once. 
+
+This module provides:
+
+* A general type-class `MonadMemo` for monads @m@ that support memoization at certain argument types @a@. 
+
+* A default trie-based implementation when @a@ is enumerable, and a curry-based implementation when @a@ is a pair type. 
+
+* A general implementation `generalmemoize` for probability, using memo-tables.
+
+* A memoized recursion combinator, `memrec`. 
+
+
+For illustrations, see the [graph example](https://lazyppl-team.github.io/GraphDemo.html), [clustering](https://lazyppl-team.github.io/ClusteringDemo.html), [additive clustering](https://lazyppl-team.github.io/AdditiveClusteringDemo.html), or the [infinite relational model](https://lazyppl-team.github.io/IrmDemo.html). 
+-}
+
+module LazyPPL.Distributions.Memoization (MonadMemo, memoize, generalmemoize, memrec) where
+
+import Control.Monad
+import Control.Monad.Extra
+import Control.Monad.State.Lazy (State, get, put, runState, state)
+import Data.IORef
+import Data.List
+import Data.Map (empty, insert, keys, lookup, size)
+import Debug.Trace
+import LazyPPL 
+import System.IO.Unsafe
+
+{-| Type class for memoizable argument types @a@ under a monad @m@ -}
+class (Monad m) => MonadMemo m a where
+  memoize :: (a -> m b) -> m (a -> b)
+  default memoize :: (Enum a) => (a -> m b) -> m (a -> b)
+  memoize f = 
+    do
+      t <- ini 0 (f . toEnum)
+      return $ \x -> look t (fromEnum x)
+
+{-- Basic trie based integer-indexed memo table.
+    NB Currently ignores negative integers --}
+data BinTree a = Branch a (BinTree a) (BinTree a)
+
+ini :: (Monad m) => Int -> (Int -> m a) -> m (BinTree a)
+ini n f = do x <- f n; l <- ini (2 * n + 1) f; r <- ini (2 * n + 2) f; return $ Branch x l r
+
+look :: BinTree a -> Int -> a
+look (Branch x l r) 0 = x
+look (Branch _ l r) n = if even n then look r (n `div` 2) else look l (n `div` 2)
+
+{-| Implementation for enumerable types using tries -}
+instance (Monad m) => MonadMemo m Int
+
+{-| Implementation for pair types using currying -}
+instance (Monad m, MonadMemo m a, MonadMemo m b) => MonadMemo m (a, b) where
+  memoize f = fmap uncurry $ memoize $ \x -> memoize $ \y -> f (x, y)
+
+{-| A general memoization method when @m@ is a probability monad.
+
+    We use unsafePerformIO to maintain
+    a table of calls that have already been made.
+    If @a@ is finite, we could just sample all values of @a@ in advance
+    and avoid unsafePerformIO. If @a@ is enumerable, we can use the trie method. 
+-}
+generalmemoize :: Ord a => (a -> Prob b) -> Prob (a -> b)
+generalmemoize f = Prob $ \(Tree _ gs) ->
+  unsafePerformIO $ do
+    ref <- newIORef Data.Map.empty
+    return $ \x -> unsafePerformIO $ do
+      m <- liftM (Data.Map.lookup x) (readIORef ref)
+      case m of
+        Just y -> return y
+        Nothing -> do
+          n <- readIORef ref
+          let y = runProb (f x) (gs !! (1 + size n))
+          modifyIORef' ref (Data.Map.insert x y)
+          return y
+
+{-- Stochastic memoization for recursive functions.
+    Applying 'memoize' to a recursively defined function only memoizes at the
+    top-level: recursive calls are calls to the non-memoized function.
+    'memrec' is an alternative implementation which resolves recursion and
+    memoization at the same time, so that recursive calls are also memoized.
+--}
+memrec :: Ord a => Show a => ((a -> b) -> (a -> Prob b)) -> Prob (a -> b)
+memrec f =
+  Prob $ \(Tree _ gs) ->
+    unsafePerformIO $ do
+      ref <- newIORef Data.Map.empty
+      let memoized_fixpoint = \x -> unsafePerformIO $ do
+            m <- liftM (Data.Map.lookup x) (readIORef ref)
+            case m of
+              Just y -> return y
+              Nothing -> do
+                n <- readIORef ref
+                let fix = f memoized_fixpoint
+                let y = runProb (fix x) (gs !! (1 + size n))
+                modifyIORef' ref (Data.Map.insert x y)
+                return y
+      return memoized_fixpoint