add shape test

Author: ahartikainen

test/data/shape.stan

@@ -0,0 +1,30 @@
+parameters {
+  real z;
+}
+model {
+  z ~ normal(0,1);
+}
+generated quantities {
+  int x = 1;
+  real a = 1;
+  real b[2];
+  real c[2,3];
+  vector[2] d;
+  vector[3] e[2];
+  matrix[2,3] f[4];
+  matrix[2,3] g[4,5];
+  for (n in 1:2) {
+    b[n] = 1;
+    d[n] = 1;
+    for (m in 1:3) {
+      c[n,m] = n;
+      e[n,m] = n;
+      for (k in 1:4) {
+        f[k, n, m] = n;
+        for (j in 1:4) {
+          g[k, j, n, m] = n;
+        }
+      }
+    }
+  }
+}

test/test_optimize.py

@@ -287,6 +287,36 @@ def test_variables_3d(self):
         self.assertTrue('frac_60' in vars_iters)
         self.assertEqual(vars_iters['frac_60'].shape, (8,))
 
+    def test_variables_shape(self):
+        stan = os.path.join(DATAFILES_PATH, 'shape.stan')
+        model = CmdStanModel(stan_file=stan)
+        no_data = {}
+        mle = model.optimize(
+            seed=1239812093,
+            algorithm='LBFGS',
+            init_alpha=0.001,
+            iter=100,
+            tol_obj=1e-12,
+            tol_rel_obj=1e4,
+            tol_grad=1e-8,
+            tol_rel_grad=1e7,
+            tol_param=1e-8,
+            history_size=5,
+        )
+        for var, shape in {
+            "x": float,
+            "a": float,
+            "b": (2,),
+            "c": (2,3),
+            "d": (2,),
+            "e":(2,3),
+            "f": (4,2,3),
+            "g":(4,5,2,3),
+        }.items():
+            if isinstance(shape, tuple):
+                assert mle.stan_variable(var).shape == shape
+            else:
+                assert isinstance(mle.stan_variable(var), shape)
 
 class OptimizeTest(unittest.TestCase):
     def test_optimize_good(self):