I’m trying to think of a way to efficiently implement a model/guide in numpyro that has two vectors sampled from a joint MVN, where the ith entries of each vector are correlated. Namely,
[\beta_1, \beta_2] ~ MVN_2k([0_k, 0_k], [[D_1, I_k r],[I_k r, D_2]]) where
D_1, D_2 are diagonal matrices,
I_k is the identity, and
r is a correlation parameter.
I implemented this in a model by manually doing an independent and conditional sample using the following code
# so just implement the conditional manually for now rho = numpyro.param("rho", 0.) sd_1 = some_func_1() # this is a vector to capture a diagonal covar D_1 sd_2 = some_func_2() # this is a vector to capture a diagonal covar D_2 beta_1 = numpyro.sample("beta_1", dist.Normal(0., sd_1)) beta_2 = numpyro.sample("beta_2", dist.Normal(rho * sd_2 / sd_1 * beta_1, jnp.sqrt(1 - rho**2) * sd_2))
What is less clear to me is how to ensure the posterior guide allows for correlation across the ith entry in
beta_2 using this approach. In theory I could create
k 2x2 MVNs, but those would not refer to the original parameters of
beta_2, which is required to map from model to guide. Any thoughts/help would be greatly appreciated.