I’m writing my own probability function, and I would like this function to be flexible with regard to the dimensionality of the data input argument. Basically, I want the function to mimic the behavior of Stan’s built-in probability functions, where “the pseudotype `reals`

is used to indicate that an argument position may be vectorized”.

In the `model`

block, it would be convenient if my function returns the summed log probability when given a 1-dimensional array of reals as input. However, in the `generated quantities`

block I would like to loop over the array to obtain the log probability of each individual input value.

As a concrete example, consider the following implementation of the inverse Gaussian (a.k.a. Wald) distribution:

```
functions {
/**
* Log density function of the inverse Gaussian (a.k.a. Wald) distribution.
*
* @param y The input value(s) whose probability needs to be evaluated.
* @param mu Mean of the inverse Gaussian distribution.
* @param lambda Shape of the inverse Gaussian distribution.
*
* @return The log of the inverse Gaussian density of y, given mean mu and shape lambda.
*
* @throws If the input value or any of the model parameters are not positive.
*
*/
real wald_lpdf(data real[] y, real mu, real lambda) {
// initialise local variables
vector[num_elements(y)] y_vector;
real term1; // constant term of log density function (does not depend on input y)
vector[num_elements(y)] term2; // term of log density function that depends on input y
vector[num_elements(y)] prob; // density of input y, given model parameters
// check that the model parameters have legal values
if (is_nan(mu) || is_inf(mu) || mu <= 0) {
reject("wald_lpdf: mu must be finite and positive; ",
"found mu = ", mu);
}
if (is_nan(lambda) || is_inf(lambda) || lambda <= 0) {
reject("wald_lpdf: lambda must be finite and positive; ",
"found lambda = ", lambda);
}
// convert the one-dimensional array y to a vector, to enable vectorised operations
y_vector = to_vector(y);
// calculate the probability of input y, given model parameters
term1 = 0.5 * log(lambda / (2 * pi()));
term2 = square(y_vector - mu) ./ (y_vector * square(mu));
prob = exp(term1 - 1.5 * log(y_vector) - 0.5 * lambda * term2);
return sum(log(prob));
}
}
data {
int N;
real<lower=0> y[N];
}
parameters {
real<lower=0> mu;
real<lower=0> lambda;
}
model {
y ~ wald_lpdf(mu, lambda);
}
generated quantities {
vector[N] log_lik;
for (i in 1:N) {
log_lik[i] = wald_lpdf(y[i] | mu, lambda); // this won't work
}
}
```

Is there a way to write the `_lpdf`

function so that it can flexibly handle both `real[] y`

as input (as above in the `model`

block) and `real y`

as input (as above in the `generated quantities`

block)?

Many thanks in advance for your suggestions!