Merge pull request #770 from null-a/guide-band-aid2

Add `noAutoGuide` option

src/guide.js

@@ -81,15 +81,14 @@ function runDistThunkCond(s, k, a, env, maybeThunk, alternate) {
       alternate(s, k, a);
 }
 
-function getDist(maybeThunk, env, s, a, k) {
+// Runs the guide thunk passed to a `sample` statement, and returns
+// the guide distribution returned by the thunk. When no guide is
+// given, then the 'noAutoGuide' flag determines whether to
+// automatically generate a suitable guide distribution or return
+// `null`.
+function getDist(maybeThunk, noAutoGuide, targetDist, env, s, a, k) {
   return runDistThunkCond(s, k, a, env, maybeThunk, function(s, k, a) {
-    return k(s, null);
-  });
-}
-
-function getDistOrAuto(maybeThunk, targetDist, env, s, a, k) {
-  return runDistThunkCond(s, k, a, env, maybeThunk, function(s, k, a) {
-    return k(s, independent(targetDist, a, env));
+    return k(s, noAutoGuide ? null : independent(targetDist, a, env));
   });
 }
 
@@ -334,6 +333,5 @@ function squishToInterval(interval) {
 module.exports = {
   independent: independent,
   runThunk: runThunk,
-  getDist: getDist,
-  getDistOrAuto: getDistOrAuto
+  getDist: getDist
 };

src/inference/elbo.ad.js

@@ -243,7 +243,10 @@ module.exports = function(env) {
 
     sample: function(s, k, a, dist, options) {
       options = options || {};
-      return guide.getDistOrAuto(options.guide, dist, env, s, a, function(s, guideDist) {
+      return guide.getDist(options.guide, options.noAutoGuide, dist, env, s, a, function(s, guideDist) {
+        if (!guideDist) {
+          throw new Error('ELBO: No guide distribution to optimize.');
+        }
 
         var ret = this.sampleGuide(guideDist, options);
         var val = ret.val;

src/inference/eubo.ad.js

@@ -115,7 +115,11 @@ module.exports = function(env) {
     sample: function(s, k, a, dist, options) {
       'use ad';
       options = options || {};
-      return guide.getDistOrAuto(options.guide, dist, env, s, a, function(s, guideDist) {
+      return guide.getDist(options.guide, options.noAutoGuide, dist, env, s, a, function(s, guideDist) {
+        if (!guideDist) {
+          throw new Error('EUBO: No guide distribution to optimize.');
+        }
+
         var rel = util.relativizeAddress(env, a);
         var guideVal = this.trace.findChoice(this.trace.baseAddress + rel).val;
         assert.notStrictEqual(guideVal, undefined);

src/inference/forwardSample.js

@@ -66,9 +66,12 @@ module.exports = function(env) {
     sample: function(s, k, a, dist, options) {
       if (this.opts.guide) {
         options = options || {};
-        return guide.getDistOrAuto(options.guide, dist, env, s, a, function(s, guideDist) {
-          return k(s, guideDist.sample());
-        });
+        return guide.getDist(
+            options.guide, options.noAutoGuide, dist, env, s, a,
+            function(s, maybeGuideDist) {
+              var d = maybeGuideDist || dist;
+              return k(s, d.sample());
+            });
       } else {
         return k(s, dist.sample());
       }

src/inference/smc.js

@@ -74,16 +74,8 @@ module.exports = function(env) {
   SMC.prototype.sample = function(s, k, a, dist, options) {
     options = options || {};
     var thunk = (this.importanceOpt === 'ignoreGuide') ? undefined : options.guide;
-    return guide.getDist(thunk, env, s, a, function(s, maybeDist) {
-
-      // maybeDist will be null if either the 'ignoreGuide' option is
-      // set, or no guide is specified in the program.
-
-      // Auto guide if requested.
-      var importanceDist =
-          !maybeDist && (this.importanceOpt === 'autoGuide') ?
-          guide.independent(dist, a, env) :
-          maybeDist;
+    var noAutoGuide = (this.importanceOpt !== 'autoGuide') || options.noAutoGuide;
+    return guide.getDist(thunk, noAutoGuide, dist, env, s, a, function(s, importanceDist) {
 
       var _val, choiceScore, importanceScore;
       if (importanceDist) {

tests/test-data/deterministic/expected/guides.json

@@ -1,6 +1,3 @@
 {
-  "result": [
-    1,
-    6
-  ]
+  "result": [true, true, true, true, true, true, true]
 }

tests/test-data/deterministic/models/guides.wppl

@@ -1,4 +1,43 @@
+var numParamsCreatedBy = function(thunk) {
+  setFreshParamsId();
+  thunk();
+  return _.size(getParams());
+};
+
 [
-  param({mu: 1, sigma: 0}),
-  T.sumreduce(param({dims: [3, 2], mu: 1, sigma: 0}))
+  param({mu: 1, sigma: 0}) === 1,
+  T.sumreduce(param({dims: [3, 2], mu: 1, sigma: 0})) === 6,
+
+  // Check (indirectly) that a guide is automatically generated, by
+  // checking that a parameter is created.
+
+  numParamsCreatedBy(function() {
+    Infer({method: 'SMC', particles: 1, importance: 'default', model() {
+      return flip();
+    }});
+  }) === 0,
+
+  numParamsCreatedBy(function() {
+    Infer({method: 'SMC', particles: 1, importance: 'autoGuide', model() {
+      return flip();
+    }});
+  }) === 1,
+
+  numParamsCreatedBy(function() {
+    Infer({method: 'SMC', particles: 1, importance: 'autoGuide', model() {
+      return sample(Bernoulli({p: 0.5}), {noAutoGuide: true});
+    }});
+  }) === 0,
+
+  numParamsCreatedBy(function() {
+    Infer({method: 'forward', samples: 1, guide: true, model() {
+      return sample(Bernoulli({p: 0.5}));
+    }});
+  }) === 1,
+
+  numParamsCreatedBy(function() {
+    Infer({method: 'forward', samples: 1, guide: true, model() {
+      return sample(Bernoulli({p: 0.5}), {noAutoGuide: true});
+    }});
+  }) === 0
 ];