VARMAX parameters had wrong shape

Author: jessegrabowski

pymc_experimental/statespace/models/VARMAX.py

@@ -220,11 +220,11 @@ def param_info(self) -> dict[str, dict[str, Any]]:
                 "constraints": "Positive Semi-definite",
             },
             "ar_params": {
-                "shape": (self.k_states, self.p, self.k_states),
+                "shape": (self.k_endog, self.p, self.k_endog),
                 "constraints": "None",
             },
             "ma_params": {
-                "shape": (self.k_states, self.q, self.k_states),
+                "shape": (self.k_endog, self.q, self.k_endog),
                 "constraints": "None",
             },
         }