WIP: Stochastic Variational Inference

rainier-sampler/src/main/scala/com/stripe/rainier/sampler/RNG.scala → rainier-base/src/main/scala/com/stripe/rainier/RNG.scala

@@ -1,7 +1,6 @@
-package com.stripe.rainier.sampler
+package com.stripe.rainier
 
 import scala.util.Random
-import Log._
 
 trait RNG {
   def standardUniform: Double
@@ -18,8 +17,6 @@ object RNG {
 }
 
 final case class ScalaRNG(seed: Long) extends RNG {
-  FINE.log("Initializing RNG with seed %d", seed)
-
   val rand: Random = new Random(seed)
   def standardUniform: Double = rand.nextDouble
   def standardNormal: Double = rand.nextGaussian

rainier-benchmark/src/main/scala/com/stripe/rainier/bench/sbc/SBCBenchmark.scala

@@ -1,11 +1,12 @@
 package com.stripe.rainier.bench.sbc
 
+import com.stripe.rainier.RNG
 import org.openjdk.jmh.annotations._
-import java.util.concurrent.TimeUnit
 
+import java.util.concurrent.TimeUnit
 import com.stripe.rainier.compute._
 import com.stripe.rainier.core._
-import com.stripe.rainier.sampler.{RNG, DensityFunction}
+import com.stripe.rainier.sampler.DensityFunction
 
 @BenchmarkMode(Array(Mode.SampleTime))
 @OutputTimeUnit(TimeUnit.MICROSECONDS)
@@ -35,7 +36,7 @@ abstract class SBCBenchmark {
 
   @Benchmark
   def run(): Unit =
-    df.update(params)
+    df.update(RNG.default, params)
 }
 
 class NormalBenchmark extends SBCBenchmark {
@@ -91,6 +92,7 @@ class BinomialPoissonApproximationBenchmark extends SBCBenchmark {
     SBC(Uniform(0, 0.04))((x: Real) => Binomial(x, 200))
 }
 
+/*
 class GaussianMixtureBenchmark extends SBCBenchmark {
   def sbc =
     SBC(Uniform(0, 1))(
@@ -102,3 +104,5 @@ class GaussianMixtureBenchmark extends SBCBenchmark {
           )
       ))
 }
+
+ */

rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/ARK.scala

@@ -1,5 +1,7 @@
 package com.stripe.rainier.bench.stan
 
+import com.stripe.rainier.compute.Real
+import com.stripe.rainier.compute.Real.parameter
 import com.stripe.rainier.core._
 
 //https://github.com/stan-dev/stat_comp_benchmarks/tree/master/benchmarks/arK
@@ -12,7 +14,7 @@ class ARK extends ModelBenchmark {
     val betas = Normal(0, 10).latentVec(5)
 
     5.until(ys.size).foldLeft(Model.empty) { (m, t) =>
-      val mu = 1.to(5).foldLeft(alpha) {
+      val mu = 1.to(5).foldLeft(alpha: Real) {
         case (mu, k) =>
           mu + betas(k - 1) * ys(t - k)
       }

rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/LowDimGaussMix.scala

@@ -6,6 +6,7 @@ import com.stripe.rainier.core._
 //https://github.com/stan-dev/stat_comp_benchmarks/tree/master/benchmarks/low_dim_gauss_mix
 //Stan: Gradient evaluation took 0.000292 seconds
 //JMH:  649.555 ± 2.572  us/op
+/*
 class LowDimGaussMix extends ModelBenchmark {
   def model = {
     val (mu1, sigma1) = muSigma()
@@ -275,3 +276,4 @@ class LowDimGaussMix extends ModelBenchmark {
     1.0084427471213, -1.6306611574291, -4.35233903314464, 3.40936081830715,
     2.75002324260943, 0.760843596839809)
 }
+ */

rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/ModelBenchmark.scala

@@ -1,10 +1,11 @@
 package com.stripe.rainier.bench.stan
 
+import com.stripe.rainier.RNG
 import org.openjdk.jmh.annotations._
-import java.util.concurrent.TimeUnit
 
+import java.util.concurrent.TimeUnit
 import com.stripe.rainier.core._
-import com.stripe.rainier.sampler.{RNG, DensityFunction}
+import com.stripe.rainier.sampler.DensityFunction
 
 @BenchmarkMode(Array(Mode.SampleTime))
 @OutputTimeUnit(TimeUnit.SECONDS)
@@ -28,7 +29,7 @@ abstract class ModelBenchmark {
 
   @Benchmark
   def run(): Unit =
-    df.update(params)
+    df.update(RNG.default, params)
 
   @Benchmark
   def build(): Unit =
@@ -46,7 +47,7 @@ object ModelBenchmarks {
   def main(args: Array[String]): Unit = {
     new ARK().main()
     new EightSchools().main()
-    new LowDimGaussMix().main()
+//    new LowDimGaussMix().main()
     new GLMMPoisson2().main()
   }
 }

rainier-compute/src/main/scala/com/stripe/rainier/compute/Compiler.scala

@@ -1,14 +1,15 @@
 package com.stripe.rainier.compute
 
-import com.stripe.rainier.ir
+import com.stripe.rainier.{RNG, ir}
 
 final case class Compiler(methodSizeLimit: Int, classSizeLimit: Int) {
   def compile(parameters: Seq[Parameter],
-              real: Real): Array[Double] => Double = {
+              real: Real): (RNG, Array[Double]) => Double = {
     val cf = compile(parameters.map(_.param), List(("base", real)))
-    return { array =>
-      val globalBuf = new Array[Double](cf.numGlobals)
-      ir.CompiledFunction.output(cf, array, globalBuf, 0)
+    return {
+      case (rng, array) =>
+        val globalBuf = new Array[Double](cf.numGlobals)
+        ir.CompiledFunction.output(cf, rng, array, globalBuf, 0)
     }
   }
   def compileTargets(group: TargetGroup): ir.DataFunction = {

rainier-compute/src/main/scala/com/stripe/rainier/compute/Evaluator.scala

@@ -33,7 +33,8 @@ class Evaluator(var cache: Map[Real, Double]) {
       Math.pow(toDouble(base), toDouble(exponent))
     case l: Lookup =>
       toDouble(l.table(toDouble(l.index).toInt - l.low))
-    case p: Parameter => sys.error(s"No value provided for $p")
+    case p: Parameter     => sys.error(s"No value provided for $p")
+    case Latent(value, _) => toDouble(value)
   }
 
   def compare(x: Real, y: Real): Int = toDouble(x).compare(toDouble(y))

rainier-compute/src/main/scala/com/stripe/rainier/compute/Gradient.scala

@@ -59,6 +59,10 @@ private object Gradient {
             //no gradient
             visit(c.left)
             visit(c.right)
+
+          case Latent(value, _) =>
+            diff(value).register(diff(real))
+            visit(value)
         }
       }
     }

rainier-compute/src/main/scala/com/stripe/rainier/compute/PartialEvaluator.scala

@@ -80,6 +80,8 @@ class PartialEvaluator(var noChange: Set[Real], rowIndex: Int) {
         (l, false)
     case p: Parameter =>
       (p, false)
+    case Latent(value, _) =>
+      apply(value)
   }
 }
 

rainier-compute/src/main/scala/com/stripe/rainier/compute/Real.scala

@@ -10,41 +10,54 @@ sealed trait Real {
   def bounds: Bounds
 
   def +(other: Real): Real = RealOps.add(this, other)
+
   def *(other: Real): Real = RealOps.multiply(this, other)
 
   def unary_- : Real = this * (-1)
+
   def -(other: Real): Real = this + (-other)
+
   def /(other: Real): Real = RealOps.divide(this, other)
 
   def min(other: Real): Real = RealOps.min(this, other)
+
   def max(other: Real): Real = RealOps.max(this, other)
 
   def pow(exponent: Real): Real = RealOps.pow(this, exponent)
 
   def exp: Real = RealOps.unary(this, ir.ExpOp)
+
   def log: Real = RealOps.unary(this, ir.LogOp)
 
   def sin: Real = RealOps.unary(this, ir.SinOp)
+
   def cos: Real = RealOps.unary(this, ir.CosOp)
+
   def tan: Real = RealOps.unary(this, ir.TanOp)
 
   def asin: Real = RealOps.unary(this, ir.AsinOp)
+
   def acos: Real = RealOps.unary(this, ir.AcosOp)
+
   def atan: Real = RealOps.unary(this, ir.AtanOp)
 
   def sinh: Real = (this.exp - (-this).exp) / 2
+
   def cosh: Real = (this.exp + (-this).exp) / 2
+
   def tanh: Real = this.sinh / this.cosh
 
   def abs: Real = RealOps.unary(this, ir.AbsOp)
 
   def logit: Real = -((Real.one / this - 1).log)
+
   def logistic: Real = Real.one / (Real.one + (-this).exp)
 }
 
 object Real {
   implicit def apply[N](value: N)(implicit toReal: ToReal[N]): Real =
     toReal(value)
+
   def seq[A](as: Seq[A])(implicit toReal: ToReal[A]): Seq[Real] =
     as.map(toReal(_))
 
@@ -61,6 +74,7 @@ object Real {
   }
 
   def parameter(): Parameter = new Parameter(new Prior(Real.zero))
+
   def parameter(fn: Parameter => Real): Parameter = {
     val x = parameter()
     x.prior = new Prior(fn(x))
@@ -78,16 +92,21 @@ object Real {
   }
 
   def doubles(seq: Seq[Double]): Real = new Column(seq.toArray)
+
   def longs(seq: Seq[Long]): Real = doubles(seq.map(_.toDouble))
 
   def eq(left: Real, right: Real, ifTrue: Real, ifFalse: Real): Real =
     lookupCompare(left, right, ifFalse, ifTrue, ifFalse)
+
   def lt(left: Real, right: Real, ifTrue: Real, ifFalse: Real): Real =
     lookupCompare(left, right, ifFalse, ifFalse, ifTrue)
+
   def gt(left: Real, right: Real, ifTrue: Real, ifFalse: Real): Real =
     lookupCompare(left, right, ifTrue, ifFalse, ifFalse)
+
   def lte(left: Real, right: Real, ifTrue: Real, ifFalse: Real): Real =
     lookupCompare(left, right, ifFalse, ifTrue, ifTrue)
+
   def gte(left: Real, right: Real, ifTrue: Real, ifFalse: Real): Real =
     lookupCompare(left, right, ifTrue, ifTrue, ifFalse)
 
@@ -110,22 +129,32 @@ object Real {
 sealed trait Constant extends Real {
   def isZero: Boolean =
     bounds.lower == 0.0 && bounds.upper == 0.0
+
   def isOne: Boolean =
     bounds.lower == 1.0 && bounds.upper == 1.0
+
   def isTwo: Boolean =
     bounds.lower == 2.0 && bounds.upper == 2.0
+
   def isPosInfinity: Boolean =
     bounds.lower.isPosInfinity && bounds.upper.isPosInfinity
+
   def isNegInfinity: Boolean =
     bounds.lower.isNegInfinity && bounds.upper.isNegInfinity
+
   def isPositive: Boolean =
     bounds.lower >= 0.0
 
   def getDouble: Double
+
   def map(fn: Double => Double): Constant
+
   def mapWith(other: Constant)(fn: (Double, Double) => Double): Constant
+
   def +(other: Constant): Constant = ConstantOps.add(this, other)
+
   def *(other: Constant): Constant = ConstantOps.multiply(this, other)
+
   def /(other: Constant): Constant = ConstantOps.divide(this, other)
 }
 
@@ -142,8 +171,11 @@ object Constant {
 
 final private case class Scalar(value: Double) extends Constant {
   val bounds = Bounds(value, value)
+
   def getDouble = value
+
   def map(fn: Double => Double) = Scalar(fn(value))
+
   def mapWith(other: Constant)(fn: (Double, Double) => Double) =
     other match {
       case Scalar(v) => Scalar(fn(value, v))
@@ -158,8 +190,11 @@ final private[rainier] class Column(val values: Array[Double])
     extends Constant {
   val param = new ir.Param
   val bounds = Bounds(values.min, values.max)
+
   def getDouble = sys.error("Not a scalar")
+
   def map(fn: Double => Double) = new Column(values.map(fn))
+
   def mapWith(other: Constant)(fn: (Double, Double) => Double) =
     other match {
       case Scalar(v) =>
@@ -186,6 +221,17 @@ final private[rainier] class Parameter(var prior: Prior) extends NonConstant {
 
 private[rainier] class Prior(val density: Real)
 
+case class FnCall(className: String,
+                  methodName: String,
+                  args: List[Real] = List.empty)
+    extends NonConstant {
+  val bounds: Bounds = Bounds(Double.NegativeInfinity, Double.PositiveInfinity)
+}
+
+case class Latent(value: Real, generator: Real) extends NonConstant {
+  val bounds = value.bounds
+}
+
 final private case class Unary(original: NonConstant, op: ir.UnaryOp)
     extends NonConstant {
   val bounds = op match {

rainier-compute/src/main/scala/com/stripe/rainier/compute/Target.scala

@@ -120,6 +120,9 @@ object TargetGroup {
           case l: Lookup =>
             loop(l.index)
             l.table.foreach(loop)
+          // NB: this should only happen for guide in variational inference
+          case Latent(value, _) =>
+            loop(value)
         }
       }
 
@@ -193,8 +196,12 @@ object TargetGroup {
 
             if (indexState.hasParameter)
               state.nonlinearOp
-            else
+            else {
               state
+            }
+          // NB: for variational inference, should stop here
+          case Latent(value, _) =>
+            loop(value)
         }
         seen += (r -> result)
         result

rainier-compute/src/main/scala/com/stripe/rainier/compute/Translator.scala

@@ -20,7 +20,10 @@ private class Translator {
           binaryExpr(toExpr(base), toExpr(exponent), PowOp)
         case Compare(left, right) =>
           binaryExpr(toExpr(left), toExpr(right), CompareOp)
-        case l: Lookup => lookupExpr(l)
+        case l: Lookup        => lookupExpr(l)
+        case Latent(value, _) => toExpr(value)
+        case FnCall(className, methodName, args) =>
+          VarDef(Sym.freshSym, FnIR(className, methodName, args.map(toExpr)))
       }
       reals += r -> expr
       expr