Make the `proposal` argument types explicit

What problem does this solve?

While reading the MCMC code I kept hitting proposal: Any, first in
mcmc_posterior.py:50.
It took me a while to figure out what a proposal has to implement, and the
answer depends on where it is used:

• The samplers only call sample() and log_prob():
  mcmc/init_strategy.py:29, :38, :68, importance/sir.py:16, rejection/rejection.py:21, :109, utils/simulation_utils.py:23.
• check_if_proposal_has_default_x (base.py:1449) also checks default_x, and only when the proposal is a NeuralPosterior.
• MCMCPosterior.to() calls proposal.to(device) (mcmc_posterior.py:162), as do the rejection and importance posteriors (rejection_posterior.py:87, importance_posterior.py:93). Neither torch.distributions.Distribution nor NeuralPosterior declares to, though the priors returned by process_prior and the concrete posteriors both have it.
• The tests pass objects that are none of the above:
  rejection_sampling_test.py:16, :59, :74 define only sample and log_prob.

The trainer side already writes the contract down as Optional[Union[Distribution, NeuralPosterior]] (_contracts.py:176, :202).

I would like to add a Proposal Protocol in sbi/sbi_types.py, next to SampleProposal and CustomPrior, that captures sample() and log_prob(), and use it for the sampler functions, with the existing union in check_if_proposal_has_default_x. CustomPrior already describes those two methods, so reusing or renaming it is fine with me if you prefer not to add another type.

One thing I could not decide on my own: MCMCPosterior.to() needs to(), which neither base class declares. Should the protocol include to, should NeuralPosterior declare to and log_prob, or should those calls stay as they are?

The pyright config disables the argument and attribute checks, so I am not proposing this for CI coverage. I just think the signatures should tell a reader what a proposal is. No runtime behavior changes.

Proposed solution

I would like to add a Proposal Protocol in sbi/sbi_types.py, next to SampleProposal and CustomPrior, that captures sample() and log_prob(), and use it for the sampler functions, with the existing union in check_if_proposal_has_default_x. CustomPrior already describes those two methods, so reusing or renaming it is fine with me if you prefer not to add another type.

class Proposal(Protocol):
    """Anything that can be drawn from and scored, e.g. a prior or a posterior."""

    def sample(self, sample_shape, /) -> Tensor: ...
    def log_prob(self, value, /) -> Tensor: ...

The parameters are positional-only so that both NeuralPosterior.log_prob(theta) and Distribution.log_prob(value) satisfy the protocol.
One thing I could not decide on my own: MCMCPosterior.to() needs to(), which neither base class declares. Should the protocol include to, should NeuralPosterior declare to and log_prob, or should those calls stay as they are?

This is type annotations only. No runtime behavior changes.

Alternatives you considered

No response

Additional context

No response

Contribution

• [ x ] I would like to work on this myself.

Start in sbi/sbi_types.py, then inspect the proposal annotations in the listed sampler, posterior, and trainer files, especially the existing CustomPrior contract and the calls to to(). Run the referenced rejection_sampling_test.py tests and review the affected signatures; done means the proposal types communicate the required interfaces without changing runtime behavior.