Repository: pyro-ppl/numpyro

Issue:
The initial value of the run function is invalid.
In the MCMC class, the folllowing initialization is succeed. It correctly starts from the initial value.
```
                kernel = numpyro.infer.NUTS(
                    model,
                    max_tree_depth=max_tree_depth,
                    init_strategy=init_to_value(values=init_params),
                    target_accept_prob=target_accept_prob,
                )
```
However, the initialization in the run function below does not reflect the initial value.
(Each value of init_params are arrays of chain numbers.)
```
                mcmc.run(
                    jax.random.PRNGKey(1),
                    y=obs,
                    init_params=init_params,
                )
```


Discussion:
The `init_params` is meant to be used for `potential_fn` (should we update the docs to better reflect this?), rather than a numpyro `model`. We can modify [this line](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/infer/hmc.py#L652) into something like
```python
new_init_params, ... = ...
init_params = new_init_params if init_params is None else init_params
```
Note that the main difference is `init_params` are samples in unconstrained (real) domain of the `model`, while `init_to_value` happens in the constrained domain of the model.
Let us know if you want us to allow using `init_params` for `model`, or to update the docs to reflect that this only takes effect for models with `potential_fn` specified.
Thank you for your detailed explanation.
I think it is better to allow using init_params for model like as init_to_value.
Do you want to make a PR for the change proposed in https://github.com/pyro-ppl/numpyro/issues/1539#issuecomment-1438987752?
Yes, I would like to.

Resolving patch:
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -649,7 +649,12 @@ def __init__(
 
     def _init_state(self, rng_key, model_args, model_kwargs, init_params):
         if self._model is not None:
-            init_params, potential_fn, postprocess_fn, model_trace = initialize_model(
+            (
+                new_init_params,
+                potential_fn,
+                postprocess_fn,
+                model_trace,
+            ) = initialize_model(
                 rng_key,
                 self._model,
                 dynamic_args=True,
@@ -658,6 +663,8 @@ def _init_state(self, rng_key, model_args, model_kwargs, init_params):
                 model_kwargs=model_kwargs,
                 forward_mode_differentiation=self._forward_mode_differentiation,
             )
+            if init_params is None:
+                init_params = new_init_params
             if self._init_fn is None:
                 self._init_fn, self._sample_fn = hmc(
                     potential_fn_gen=potential_fn,
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -515,7 +515,9 @@ def warmup(
         :param bool collect_warmup: Whether to collect samples from the warmup phase. Defaults
             to `False`.
         :param init_params: Initial parameters to begin sampling. The type must be consistent
-            with the input type to `potential_fn`.
+            with the input type to `potential_fn` provided to the kernel. If the kernel is
+            instantiated by a numpyro model, the initial parameters here correspond to latent
+            values in unconstrained space.
         :param kwargs: Keyword arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init`
             method. These are typically the keyword arguments needed by the `model`.
         """
@@ -546,7 +548,9 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
             `"adapt_state.step_size"` can be used to collect step sizes at each step.
         :type extra_fields: tuple or list of str
         :param init_params: Initial parameters to begin sampling. The type must be consistent
-            with the input type to `potential_fn`.
+            with the input type to `potential_fn` provided to the kernel. If the kernel is
+            instantiated by a numpyro model, the initial parameters here correspond to latent
+            values in unconstrained space.
         :param kwargs: Keyword arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init`
             method. These are typically the keyword arguments needed by the `model`.
 


Test patch: