I am working on pystan==3.6.0, Python 3.8.11

I have a list of numpy.ndarray (tensors) called tensor_values (which can be of variable shape, However, all the tensors in the list will be of same shape - I will show the range of shapes at the end) and a list of float values of weights called weights_norm_1

This is my input data

```
if isinstance(tensor_values[0], np.ndarray):
max_value = len(tensor_values[0]) #Shape_of_one_of_input_tensor to give stan
else:
max_value = 1
stan_data = {
'N': len(tensor_values),
'n': len(weights_norm_1),
'P': max_value, #Shape_of_one_of_input_tensor
'p': [t.flatten() if t.ndim > 0 else np.array([t]) for t in tensor_values],
'lambda': 1,
'mu_c': 1,
'weights': np.array(weights_norm_1)
}
```

My Stan code is as follows

```
data {
int<lower=1> N; // Number of input values
int<lower=1> n; // length of weight_norm_1
int<lower=1> P; // Number of parameters in each tensor
vector[P] p[N]; // Parameters for each input tensor
real<lower=0> lambda; // Variance of each input tensor's parameters
real mu_c; // Mean of the bias for each input tensor
vector[n] weights; // Weights for each input tensor, computed externally weight_norm_1
}
parameters {
vector[P] p_m; // Global mean parameter
vector[N] b_c; // Bias for each input tensor
}
model {
vector[n] sigma_c; // Standard deviations for input tensor biases
// Priors
p_m ~ normal(0, 1); // Prior on the global mean parameters
// Compute sigma_c based on weights
for (i in 1:n) {
sigma_c[i] = 1 - weights[i];
}
b_c ~ normal(mu_c, sigma_c); // Prior on biases
// Likelihood
for (i in 1:N) {
p[i] ~ normal(p_m * b_c[i], lambda); // Model for each input tensors parameters
}
}
```

Python code that handles the output of stan

```
posterior = stan.build(model_code, data=stan_data)
fit = posterior.sample(num_chains=4, num_samples=1000)
print(f" fit.summary: {fit}")
```

I want to calculate bayesian mean, the output should have the same dimensions as of the list of tensor_values. How can I fix the code to be able to handle the input and output as required?

Following is the list of all the potential input tensor shapes

```
Shape of element: ()
Shape of element: (1,)
Shape of element: (128,)
Shape of element: (128, 128, 3, 3, 3)
Shape of element: (128, 256, 1, 1, 1)
Shape of element: (128, 256, 3, 3, 3)
Shape of element: (128, 64, 3, 3, 3)
Shape of element: (256,)
Shape of element: (256, 128, 3, 3, 3)
Shape of element: (256, 256, 3, 3, 3)
Shape of element: (256, 512, 1, 1, 1)
Shape of element: (256, 512, 3, 3, 3)
Shape of element: (32,)
Shape of element: (32, 32, 3, 3, 3)
Shape of element: (32, 4, 3, 3, 3)
Shape of element: (32, 64, 1, 1, 1)
Shape of element: (4,)
Shape of element: (4, 64, 3, 3, 3)
Shape of element: (512,)
Shape of element: (512, 256, 3, 3, 3)
Shape of element: (512, 512, 3, 3, 3)
Shape of element: (64,)
Shape of element: (64, 128, 1, 1, 1)
Shape of element: (64, 128, 3, 3, 3)
Shape of element: (64, 32, 3, 3, 3)
Shape of element: (64, 64, 3, 3, 3)
```