Oh no, I had two notebooks open last night and was looking at the wrong one when making some of those edits.
For anyone else who finds this, I did have to make a change to the step sampler. The astep function should be implemented as follows:
def astep(self, q0: RaveledVars) -> RaveledVars:
point_Amap_info = q0.point_map_info
q0 = q0.data
q = self.proposal_dist(size=self.ndim)
return RaveledVars(q, point_map_info), []
So it should return a tuple of the RaveledVars and an empty list. The list is expected by the parent class ArrayStep and should contain a dictionary of “stats” related to the step (for example related to tuning). Here, we don’t have anything that we want to keep track of, so it’s left empty