NaN occurred in optimization with ADVI

Usually, if you are seeing the the NaN problem in the first iteration, there is problem of the approximation set up that makes the score function (target to be optimized) invalid.
So, I would follow:

  1. check if the original model is set up correctly, see eg here.
  2. print the test value of the deterministic node to make sure that the value is valid, eg doing p.tag.test_value in your case.
  3. make sure the starting value of the approximation is valid:
point = apprx.groups[0].bij.rmap(apprx.params[0].eval())
point #check to see there is no NaN or Inf
for var in logistic_model.free_RVs:
    print(var.name, var.logp(point))

If all the steps above pass fine, there is a problem with setting up the approximation score function, which is not that easy to diagnose. I would put everything in a jupyter notebook, and trace the error in %debug mode.

1 Like