GSoC #9: The Final Push

Although I still have two weeks left, I’ll be at the ODYSSEY Alignment Conference next week so I’ll have limited time, which means I’ll be leading the push this week to get the main points sorted.

Last week I finished up my main PR implementing INLA end-to-end, which is just waiting on approval and should be ready to go (two unusual test cases are failing, but I believe they’re unrelated to my PR).

Essentially, the main outstanding bits are:

• Implementing an unmarginalization scheme on the latents. This will involve digging through sample_posterior_predictive, recover_marginals and a few other places here and there to try and understand how to best approach it. Either we build a FunctionGraph mapping \(\theta \rightarrow x\) and feed the idata draws through directly, or we do something along the lines of sample_posterior_predictive. I hope that this should be a fairly quick endeavour with the help of Rob and the others, as this is a (somewhat) standard procedure.

• Writing an example usage notebook. This should be a fairly quick and easy process. I’m thinking of using ICAR as the example case, and simply comparing pm.sample to pmx.fit(method="INLA") on this more nontrivial model where INLA is (theoretically) faster. However, I’m unsure to what extent I can work on this while the Jacobian minimize stuff is being patched by Ricardo and Jesse. For the moment, I can probably just use Metropolis sampling.

• Diagnose the (apparent) minimize slowdown - is there something limiting our runtime or is it simply a matter of INLA outscaling direct sampling? Again, this may be blocked by the minimize issue, but simply benchmarking both cases on Metropolis is a decent starting point. However, it is important to note that Metropolis doesn’t use gradients, and so there is no observed speedup from using adjoint methods on this benchmark.