I would like to perform HMC inference of following probability function:

where

- \mathbf{y}_j \in \mathbb{R}^{N}
- \mathbf{x}_n \in \mathbb{R}^{D}
- \mathbf{O} \in \mathbb{R}^{J \times N} \implies \mathbf{O}_j \in \mathbb{R}^N
- K: (\mathbf{X}, \mathbf{X}^{\top}) \mapsto \mathbb{R}^{N \times N} (some kernel)

In Python, I would write something like

```
def log_posterior(X):
K = kernel(X)
Q = 0
for n in range(N):
Q += multivariate_normal.logpdf(X[n], mu_x, Cov_x)
for j in range(J):
Q += multivariate_normal.logpdf(Y[j], mu_z, O[j] + K)
return Q
```

However, I want to try to optimize this (find \mathbf{X}^{\star}) using Stan. However, I am having trouble. This is what I have tried so far:

```
data {
int<lower=0> N;
int<lower=0> J;
int<lower=0> D;
vector[N] mu_y;
vector[N] Y[J];
vector[N] O[J];
real alpha;
real rho;
vector[D] mu_x;
matrix[D, D] Cov_x;
}
parameters {
vector[N] X[D];
}
transformed parameters {
matrix[N, N] Cov[J];
matrix[N, N] K = cov_exp_quad(X, X, alpha, rho);
for (j in 1:J)
Cov[j] = diag_matrix(O[j]) + K;
}
model {
X ~ multi_normal(mu_x, Cov_x);
for (j in 1:J)
Y[j] ~ multi_normal(mu_y, Cov[j]);
}
```

I get lots of various errors about dimensions being wrong, but more generally, I think my mental model of how to write in Stan is still incomplete. I’m particularly confused as to how to handle indexing across J. Any help is appreciated.