First Pystan Poisson Model

Modeling
6

I think I see a potential problem here:

transformed parameters{
    real mu[N];
    for (i in 1:N) {
        mu[i] <- a*driver_experience[i] + b;
    }
}

mu is a parameter. That means that it isn’t a temporary value. Rather, PyStan has to have enough memory to store not only a single instance of mu, but however many mu's are generated, i.e. the number of iterations times the number of parallel chains. Now if mu were just a single real value, 8 bytes in size, that wouldn’t be that big a deal. For 1000 iterations and 2 chains (which you had in your first model), that would be 8\ \mathrm{bytes} \times 1000 \times 2 or about 16 kilobytes – a fairly trivial amount of memory on a modern computer.

However, I remember that in your first model that you set N to about 1 million. So now, at each iteration mu is about 8 megabytes, and if you have the same number of iterations and chains as before, you now need to store about 16 gigabytes in total. That’s far less trivial.

I would recommend getting rid of your transformed parameters block and writing your model block like this:

model {
    a ~ normal(0, 1);  
    b ~ normal(0,5);
    y1 ~ poisson(a*driver_experience + b);   
}

That should use far less memory.