Repository: pyro-ppl/numpyro

Issue:
Inverse bjiector transformation (from constrained to unconstrained space)
Hi,
Is it possible, with the current interface, to design a function that takes as input the pytree of _constrained_ parameter values and return _unconstrained_ values, given a model with Pyro primitives? For instance, I am looking for a function that does the inverse transform of `constrain_fn()` from `numpyro.infer.util`.

Sorry to ask here but I could not find the answer in the code. So maybe there is no easy implementation of it yet.


Discussion:
I just found a way to achieve this, using [this piece of code](https://github.com/pyro-ppl/numpyro/blob/9a43b19faf547c7b7259d0e5550d1d9125e9a881/numpyro/infer/util.py#L343) embedded in `infer.util`.

If you find it useful, I can write a PR for adding a function `unconstrained_fn()` to the module, following the same interface as the existing `constrained_fn()`.
yes, please make a PR if it is helpful for your usage case. Thanks!
Ok @fehiepsi ! I will propose a PR once I am sure that I fully understand everything regarding my own test case, and I can write functions that are general enough to be useful to anyone.
Take your time, @aymgal. Please let me know if you need some clarification in the code. This feature is also useful for https://forum.pyro.ai/t/initialize-mcmc-chains-from-multiple-predetermined-starting-points/5062/10
@fehiepsi I am in the process of opening a PR. Beforehand I just wanted to check with you about the available functions in `infer.util`.
In my use cases I have a model with both `param` and `sample` sites, so I need an `constrain_fn` (and a new `unconstrain_fn`) that both supports `param`. If I am correct, the current implementation of `constrain_fn` only considers `sample` sites, leaving untouched the `param` although they can be subject to specific constraints.

I thought a simple option is to implement the following function 
```python
def get_transforms(model, model_args, model_kwargs, params):
    """
    (EXPERIMENTAL INTERFACE) Retrieve transforms functions given a NumPyro model.
    This function supports 'param' sites.

    :param model: a callable containing NumPyro primitives.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param dict params: dictionary of unconstrained values keyed by site
        names.
    :return: `dict` of transformation keyed by site names.
    """
    substituted_model = handlers.substitute(model, params)
    model_trace = handlers.trace(substituted_model).get_trace(*model_args, **model_kwargs)
    transforms = {}
    for k, v in model_trace.items():
        if v["type"] == "param":
            constraint = v["kwargs"].pop("constraint", constraints.real)
            with helpful_support_errors(v):
                transforms[k] = transforms.biject_to(constraint)
        elif (
            v["type"] == "sample"
            and not v["is_observed"]
            and not v["fn"].support.is_discrete
        ):
            with helpful_support_errors(v):
                transforms[k] = transforms.biject_to(v["fn"].support)
    return transforms
```

and then use the already existing `transform_fn` to apply either the forward (`invert=False`, from unconstrained to constrained space) or backward (`invert=True`, from constrained to unconstrained space) transform to a dict of site values. What do you think?
However, this implementation would lead to an alternate version of `constrain_fn`, which may not necessarily always lead to the same results, so I am not sure what is the best here.

I am happy to give further details if something is unclear.
> which may not necessarily always lead to the same results

Your proposal looks great to me. Why this approach might lead to a different result?
@fehiepsi One difference would be that the current `constrain_fn` and my version will not handle `param` sites the same way, so both functions won't return the same values if the numpyro model contains param primitives.
Additionally, what about function naming? Should I choose a different name than `constrain_fn`? Or maybe I can add an extra arguments to the function (e.g. `include_param={True, False}`) to bridge the gap between the two implementations?
I think `constrain_fn` only transforms sites in the inputs. If the inputs only contain sample sites then it should only return the corresponding output for those sample sites.

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -181,6 +181,10 @@ def substitute_fn(site):
             if site["type"] == "sample":
                 with helpful_support_errors(site):
                     return biject_to(site["fn"].support)(params[site["name"]])
+            elif site["type"] == "param":
+                constraint = site["kwargs"].pop("constraint", constraints.real)
+                with helpful_support_errors(site):
+                    return biject_to(constraint)(params[site["name"]])
             else:
                 return params[site["name"]]
 
@@ -193,6 +197,42 @@ def substitute_fn(site):
     }
 
 
+def get_transforms(model, model_args, model_kwargs, params):
+    """
+    (EXPERIMENTAL INTERFACE) Retrieve (inverse) transforms via biject_to()
+    given a NumPyro model. This function supports 'param' sites.
+    NB: Parameter values are only used to retrieve the model trace.
+
+    :param model: a callable containing NumPyro primitives.
+    :param tuple model_args: args provided to the model.
+    :param dict model_kwargs: kwargs provided to the model.
+    :param dict params: dictionary of values keyed by site names.
+    :return: `dict` of transformation keyed by site names.
+    """
+    substituted_model = substitute(model, data=params)
+    transforms, _, _, _ = _get_model_transforms(
+        substituted_model, model_args, model_kwargs
+    )
+    return transforms
+
+
+def unconstrain_fn(model, model_args, model_kwargs, params):
+    """
+    (EXPERIMENTAL INTERFACE) Given a NumPyro model and a dict of parameters,
+    this function applies the right transformation to convert parameter values
+    from constrained space to unconstrained space.
+
+    :param model: a callable containing NumPyro primitives.
+    :param tuple model_args: args provided to the model.
+    :param dict model_kwargs: kwargs provided to the model.
+    :param dict params: dictionary of constrained values keyed by site
+        names.
+    :return: `dict` of transformation keyed by site names.
+    """
+    transforms = get_transforms(model, model_args, model_kwargs, params)
+    return transform_fn(transforms, params, invert=True)
+
+
 def _unconstrain_reparam(params, site):
     name = site["name"]
     if name in params:
@@ -449,6 +489,10 @@ def _get_model_transforms(model, model_args=(), model_kwargs=None):
                 for arg in args:
                     if not isinstance(getattr(support, arg), (int, float)):
                         replay_model = True
+        elif v["type"] == "param":
+            constraint = v["kwargs"].pop("constraint", constraints.real)
+            with helpful_support_errors(v, raise_warnings=True):
+                inv_transforms[k] = biject_to(constraint)
         elif v["type"] == "deterministic":
             replay_model = True
     return inv_transforms, replay_model, has_enumerate_support, model_trace


Test patch:
diff --git a/test/infer/test_infer_util.py b/test/infer/test_infer_util.py
--- a/test/infer/test_infer_util.py
+++ b/test/infer/test_infer_util.py
@@ -32,6 +32,7 @@
     log_likelihood,
     potential_energy,
     transform_fn,
+    unconstrain_fn,
 )
 import numpyro.optim as optim
 
@@ -220,6 +221,52 @@ def model():
     assert_allclose(actual_potential_energy, expected_potential_energy)
 
 
+def test_constrain_unconstrain():
+    x_prior = dist.HalfNormal(2)
+    y_prior = dist.LogNormal(scale=3.0)  # transformed distribution
+    z_constraint = constraints.positive
+
+    def model():
+        numpyro.sample("x", x_prior)
+        numpyro.sample("y", y_prior)
+        numpyro.param("z", init_value=2.0, constraint=z_constraint)
+
+    params = {"x": jnp.array(-5.0), "y": jnp.array(7.0), "z": jnp.array(3.0)}
+    model = handlers.seed(model, random.PRNGKey(0))
+    inv_transforms = {
+        "x": biject_to(x_prior.support),
+        "y": biject_to(y_prior.support),
+        "z": biject_to(z_constraint),
+    }
+    expected_constrained_samples = partial(transform_fn, inv_transforms)(params)
+    transforms = {
+        "x": biject_to(x_prior.support).inv,
+        "y": biject_to(y_prior.support).inv,
+        "z": biject_to(z_constraint).inv,
+    }
+    expected_unconstrained_samples = partial(transform_fn, transforms)(
+        expected_constrained_samples
+    )
+
+    actual_constrained_samples = constrain_fn(model, (), {}, params)
+    actual_unconstrained_samples = unconstrain_fn(
+        model, (), {}, actual_constrained_samples
+    )
+
+    assert_allclose(expected_constrained_samples["x"], actual_constrained_samples["x"])
+    assert_allclose(expected_constrained_samples["y"], actual_constrained_samples["y"])
+    assert_allclose(expected_constrained_samples["z"], actual_constrained_samples["z"])
+    assert_allclose(
+        expected_unconstrained_samples["x"], actual_unconstrained_samples["x"]
+    )
+    assert_allclose(
+        expected_unconstrained_samples["y"], actual_unconstrained_samples["y"]
+    )
+    assert_allclose(
+        expected_unconstrained_samples["z"], actual_unconstrained_samples["z"]
+    )
+
+
 def test_model_with_mask_false():
     def model():
         x = numpyro.sample("x", dist.Normal())