That looks neat. I wonder if that should just sample and then return an InferenceData object to be consistent with the rest.