Clarify how to sample the centered parameter in LocScaleReparam

[Madhav-Kanda]

The documentation for LocScaleReparam in case of centeredness mentions that: " If None (default) learn a per-site per-element centering parameter in [0,1]", but upon looking through the implementation of the algorithm, I did not find anything relevant to finding the correct centeredness (as far as I understand the implementation by default it takes the centeredness to be 0.5 and does not optimize it).

Example:
In the case of the eight schools' examples, the ideal centeredness should be close to 0. Still, upon using LocScaleReparam with None as its parameter, we get the centeredness is 0.5, as seen below (which I believe is hard coded in the implementation of LocScaleReparam itself for the case of None).

J = 8
y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])

def eight_schools_noncentered(J, sigma, lambd, y=None):
    mu = numpyro.sample('mu', dist.Normal(2, 5))
    tau = numpyro.sample('tau', dist.HalfCauchy(5))
    with numpyro.plate('J', J):
        with numpyro.handlers.reparam(config={'theta': LocScaleReparam(centered=lambd)}):
            theta = numpyro.sample('theta', dist.Normal(mu, tau))
        numpyro.sample('obs', dist.Normal(theta, sigma), obs=y)

nuts_kernel = NUTS(eight_schools_noncentered)
mcmc = MCMC(nuts_kernel, num_warmup=1000, num_samples=1000)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, J, sigma, lambd = None,y=y, extra_fields=('potential_energy',))
mcmc.print_summary()

mcmc = MCMC(nuts_kernel, num_warmup=1000, num_samples=1000)
mcmc.run(rng_key, J, sigma, lambd = 0.5,y=y, extra_fields=('potential_energy',))
mcmc.print_summary()

In the above case both of them provide the same print_summary. Thus, I believe that either the code for the LocScaleReparam should be updated to include SVI for finding the best centeredness or the documentation should be updated accordingly.

[fehiepsi]

does not optimize it

You are running MCMC, rather than optimizing. As mentioned in the forum thread, you can use lift handler to cast the param statement to a sample statement.

[fehiepsi]

Please feel free to submit a PR to clarify the doc. :)

[Madhav-Kanda]

Sure, I will do that.