Gradients cannot simultaneously move f and x in f(x).

This applies to GPs, for example, when f is a GP with a randomly parametrized covariance kernel, and x is a randomly chosen input to the GP. This would also apply to, e.g., NIG-Normal SPs, if we supported gradients on them at all.

Single-site gradient ascent with grad_ascent(default, one, ...) works for these cases because it proposes only changing one of f or x at a time before proposing changing the other one. But grad_ascent(default, all, ...) tries to change both simultaneously, and fails.

To fix this, we must:

[ ] Change the signature of gradientOfLogDensity to include the gradient with respect to the SP's internal state in the return value.
[ ] Change the signature of gradientOfLogDensityOfData to include the gradient with respect to the SP's internal state in the return value.
[ ] Change the signature of gradientOfSimulate to include the gradient with respect to the SP's internal state in the return value.
[ ] Teach detach/regen to record the gradient with respect to the operator's internal state at each application, in addition to the gradients with respect to all the operands.

Test case with GPs:

@stochasticTest
def test_arglebargle(seed):
  ripl = get_ripl(seed=seed)
  ripl.assume('gs_expon_1',
    '(lambda () (- 0. (log_logistic (log_odds_uniform))))')
  ripl.assume('mu_0', '(normal 0 1)')
  ripl.assume('s', '(gs_expon_1)')
  ripl.assume('l', '(gs_expon_1)')
  ripl.assume('mean', '(gp_mean_const mu_0)')
  ripl.assume('cov', '(gp_cov_scale (* s s) (gp_cov_se (* l l)))')
  ripl.assume('gp', '(make_gp mean cov)')
  ripl.assume('foo', '(lookup (gp (array 1)) 0)')
  ripl.observe('(normal foo 1)', 1.2)
  ripl.infer('(grad_ascent default all 0.1 10 10)')
  ripl.sample('(gp (array 2 3))')

Test case with NIG-Normal:

@stochasticTest
def test_nig_normal_gradient(seed):
  ripl = get_ripl(seed=seed)
  ripl.assume('gs_expon_1',
    '(lambda () (- 0. (log_logistic (log_odds_uniform))))')
  ripl.assume('m', '(gs_expon_1)')
  ripl.assume('V', '(gs_expon_1)')
  ripl.assume('a', '(gs_expon_1)')
  ripl.assume('b', '(gs_expon_1)')
  ripl.assume('nig_normal', '(make_nig_normal m V a b)')
  ripl.observe('(normal (nig_normal) 1)', '0')
  ripl.infer('(grad_ascent default all 0.1 10 10)')

Failure symptom:

(run (grad_ascent default all 0.1 10 10))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Caused by

**************************************************
Stack trace from worker:
**************************************************
Traceback (most recent call last):
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/multiprocess.py", line 104, in wrapped
    res = f(*args, **kwargs)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/multiprocess.py", line 342, in <lambda>
    return safely(lambda *args,**kwargs: getattr(self.obj, attrname)(*args, **kwargs))
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/trace.py", line 500, in primitive_infer
    return infer.primitive_infer(self, exp)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/dispatch.py", line 251, in primitive_infer
    return transloop(trace, transitions, scaffolder_loop(scaffolders, doit))
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/dispatch.py", line 108, in transloop
    ct += operate()
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/dispatch.py", line 122, in doit
    ct += operate(scaffolder)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/dispatch.py", line 250, in doit
    return mixMH(trace, scaffolder, GradientAscentOperator(rate, int(steps)))
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/mh.py", line 75, in mixMH
    proposedTrace, logAlpha = operator.propose(trace, index)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/map_gradient.py", line 45, in propose
    proposed_values = self.evolve(grad, currentValues, start_grad)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/map_gradient.py", line 56, in evolve
    dxs = grad(xs)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/hmc.py", line 62, in __call__
    self.fixed_regen(values)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/hmc.py", line 72, in fixed_regen
    return self.regen(values)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/hmc.py", line 75, in regen
    registerDeterministicLKernels(self.trace, self.scaffold, self.pnodes, values)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/infer/mh.py", line 39, in registerDeterministicLKernels
    DeterministicLKernel(trace.pspAt(pnode), currentValue)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/trace.py", line 238, in pspAt
    def pspAt(self, node): return node.relevantPSP(self.spAt(node))
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/trace.py", line 232, in spAt
    def spAt(self, node): return self.madeSPAt(self.spRefAt(node).makerNode)
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/trace.py", line 256, in madeSPAt
    def madeSPAt(self, node): return self.madeSPRecordAt(node).sp
  File "/home/riastradh/venture/master/build/lib.linux-x86_64-2.7/venture/lite/trace.py", line 250, in madeSPRecordAt
    assert node.madeSPRecord is not None
AssertionError

probcomp / Venturecxx

Gradients cannot simultaneously move f and x in f(x). #632