I have a high dimensional bayesian numpyro model, where I have two sets of parameters I’m trying to estimate “x” and “z”. Each are vectors of the same length ~10,000. I’m fitting using NUTS and MCMC.

I know that “x” and “z” are highly correlated, **but only for matching pairs**. That is, I know that x_1 is correlated with z_1, and x_2 with z_2, but not x_1 with z_2 or x_2 with z_1. Put another way, I know x_i is correlated with z_j for i==j only.

Because the dimensionality is so high, if I just specify `dense_mass = [("x","z")]`

crashes due to insufficient memory.

So my question is, how do I tell NUTS to estimate only the matching pairs from “x” and “z” in the mass matrix?

Here is a basic version of my model:

```
def model(a,b,y):
x = numpyro.sample("x",dist.Normal(0,1).expand([10000]))
z = numpyro.sample("z",dist.Normal(0,1).expand([10000]))
lambda_ = jnp.exp(x*a+z*b) + 1e-16
numpyro.sample('T',dist.Poisson(lambda_),obs=y)
```