diff --git a/R/inst/stan/prophet.stan b/R/inst/stan/prophet.stan
index 7586125..d8797e7 100644
--- a/R/inst/stan/prophet.stan
+++ b/R/inst/stan/prophet.stan
@@ -98,8 +98,9 @@ data {
 }
 
 transformed data {
-  matrix[T, S] A;
-  A = get_changepoint_matrix(t, t_change, T, S);
+  matrix[T, S] A = get_changepoint_matrix(t, t_change, T, S);
+  matrix[T, K] X_sa = X .* rep_matrix(s_a', T);
+  matrix[T, K] X_sm = X .* rep_matrix(s_m', T);
 }
 
 parameters {
@@ -130,10 +131,10 @@ model {
   beta ~ normal(0, sigmas);
 
   // Likelihood
-  y ~ normal(
-    trend
-      .* (1 + X * (beta .* s_m))
-      + X * (beta .* s_a),
-      sigma_obs
-    );
+  y ~ normal_id_glm(
+    X_sa,
+    trend .* (1 + X_sm * beta),
+    beta,
+    sigma_obs
+  );
 }
diff --git a/python/stan/prophet.stan b/python/stan/prophet.stan
index cb18633..05bcb34 100644
--- a/python/stan/prophet.stan
+++ b/python/stan/prophet.stan
@@ -101,8 +101,9 @@ data {
 }
 
 transformed data {
-  matrix[T, S] A;
-  A = get_changepoint_matrix(t, t_change, T, S);
+  matrix[T, S] A = get_changepoint_matrix(t, t_change, T, S);
+  matrix[T, K] X_sa = X .* rep_matrix(s_a', T);
+  matrix[T, K] X_sm = X .* rep_matrix(s_m', T);
 }
 
 parameters {
@@ -133,10 +134,10 @@ model {
   beta ~ normal(0, sigmas);
 
   // Likelihood
-  y ~ normal(
-  trend
-  .* (1 + X * (beta .* s_m))
-  + X * (beta .* s_a),
-  sigma_obs
+  y ~ normal_id_glm(
+    X_sa,
+    trend .* (1 + X_sm * beta),
+    beta,
+    sigma_obs
   );
 }