How can I set a special custom LKJ prior?

Hi devs and community

I want to set LKJ prior for a 3x3 correlation matrix. I expect strong correlation between neighbors, that is, (0, 1) and (1, 2) and weak correlation between non-neighbors, that is, (1, 3). Is there a way to set a prior like this?

Right now I can only set concentration parameter to be either between 0 and 1 (expecting strong correlation) or greater than 1 (expecting weak correlation)

If you’re familiar with [partial correlation] (Partial correlation - Wikipedia), I guess you can modify the base Beta distribution here and define your custom LKJ as a transformed distribution of the signed stick breaking transform there.

Alternatively, it might be easier to do something like this

corr = numpyro.sample("corr", LKJ(...))
# add a little bias
bias = jnp.array([[1, 1, 0], [1, 1, 1], [0, 1, 1]])
numpyro.factor("corr_bias", dist.Normal(bias, 0.1).log_prob(corr))