Modifying cartesian product to allow for >2D input arrays

pymc3/math.py

@@ -101,11 +101,17 @@ def cartesian(*arrays):
 
     Parameters
     ----------
-    arrays: 1D array-like
-            1D arrays where earlier arrays loop more slowly than later ones
+    arrays: N-D array-like
+            N-D arrays where earlier arrays loop more slowly than later ones
     """
     N = len(arrays)
-    return np.stack(np.meshgrid(*arrays, indexing="ij"), -1).reshape(-1, N)
+    arrays_np = [np.asarray(x) for x in arrays]
+    arrays_2d = [x[:, None] if np.asarray(x).ndim == 1 else x for x in arrays_np]
+    arrays_integer = [np.arange(len(x)) for x in arrays_2d]
+    product_integers = np.stack(np.meshgrid(*arrays_integer, indexing="ij"), -1).reshape(-1, N)
+    return np.concatenate(
+        [array[product_integers[:, i]] for i, array in enumerate(arrays_2d)], axis=-1
+    )
 
 
 def kron_matrix_op(krons, m, op):

pymc3/tests/test_math.py

@@ -74,6 +74,23 @@ def test_cartesian():
     np.testing.assert_array_almost_equal(manual_cartesian, auto_cart)
 
 
+def test_cartesian_2d():
+    np.random.seed(1)
+    a = [[1, 2], [3, 4]]
+    b = [5, 6]
+    c = [0]
+    manual_cartesian = np.array(
+        [
+            [1, 2, 5, 0],
+            [1, 2, 6, 0],
+            [3, 4, 5, 0],
+            [3, 4, 6, 0],
+        ]
+    )
+    auto_cart = cartesian(a, b, c)
+    np.testing.assert_array_almost_equal(manual_cartesian, auto_cart)
+
+
 def test_kron_dot():
     np.random.seed(1)
     # Create random matrices