Expose initialization strategy in HMC

[fehiepsi]

Addresses part of #2415.

There is one issue is: if users want to change keywords of init strategies, they have to use partial, e.g. partial(init_to_median, num_samples=20). I find that it is a bit inconvenient for users. In NumPyro, users can set init_strategy= init_to_median(20) or init_to_uniform(3), init_to_value({"x": 10}), init_to_feasible().

[fritzo]

Re: usability, you could use the standard poutine trick:

def init_to_uniform(site=None, radius=2):
    if site is None:
        return functools.partial(init_to_uniform, radius=2)
    ...

Then users can pass init_to_uniform(radius=0.5) as their function.

[fritzo]

This is great!

Do I understand correctly that this allows me to specify a subset of init_params to do so in the constrained space (unlike the initial_params kwarg, which if so we should mark DEPRECATED).

[fehiepsi]

@fritzo initial_params is originally useful when users use potential_fn instead of model.

Re: usability, you could use the standard poutine trick:

Thanks, yeah we should do it!

[fehiepsi]

Do I understand correctly that this allows me to specify a subset of init_params to do so in the constrained space

Yes, you just need to specify a subset. The remaining params use the uniform strategy (default in MCMC).

[fritzo]

LGTM Thanks for implementing this!

[neerajprad]

@fritzo - does this address your concern of initialization code being triggered at various points in the flow despite the user specifying initial params?

[fehiepsi]

@neerajprad This does not. I think we should address it.

[fehiepsi]

I am fixing some typos. I'll address that too.

[neerajprad]

Thanks for clarifying, @fehiepsi. That might involve some refactoring, look forward to seeing your PR. Let me know if I can help.

[fritzo]

does this address your concern of initialization code being triggered

Yes this addresses my concern with duplicated work. The problem was, the model was running 100 random initializations even if I provided a single good one. After This PR, if I provide a single good initialization through init_to_value, no random initializations are sampled.