Return sampler state in pm.sample

[junpenglao]

Currently, pm.sample does not return step_method, which makes restart sampling impossible.
Some workaround we suggest user is to reinitialized the step_method manually using the posterior. This is mostly fine but we are not exactly restarting from the same state.
There are a few ways I see we can archive this:

• store additional sampler information in diagnostics (e.g., the mass matrix)
• return the step_method (might be difficult as sampling use Python generator)

    ...
    idata, step_methods = pm.sample(..., return_step_method_class=True)

• refactor the API to be more stateless (similar to design in TFP and BlackJAX)

[michaelosthege]

Technically we could include it in sampler stats, but that might eat a lot of bytes if we're only interested in the final state after sampling.

I don't think we should return it separately: We'll probably want to store it in the InferenceData anyway.
What is the data structure of the sampler state?
Instead of including it via sampler stats it could become its own field in the RunMeta data.

[ricardoV94]

I think this is a good motivation to refactor the step samplers. For instance to move away from global seeding:

#5093

#5378

Agree with making sampling more stateless

[HarshKhilawala]

Dear @ricardoV94 ,

I am a newcomer to PyMC and I am eager to contribute to the community. I am interested in the GSoC 2023 idea - Project: Refactor step sampler API, but before I begin working on my project proposal, I would like to gain a better understanding of the PyMC codebase and how it functions. As a student developer, I am highly motivated to learn and grow in the tech industry. I was wondering if you could provide me with any guidance on how to get started with contributing to PyMC. Your help would be greatly appreciated.

[reshamas]

Here is an issue that @ricardoV94 has pinned on getting started contributing:
#6518