From 47b36586676a0af6505f376afff2193ab5a5ed84 Mon Sep 17 00:00:00 2001
From: Christopher Krapu <ckrapu@gmail.com>
Date: Fri, 26 Feb 2021 16:16:15 -0500
Subject: [PATCH] Modifying cartesian product to allow for >2D input arrays
 (#4482)

* Modifying cartesian product to allow for more >2D input arrays
* Assert for equality in cartesian test
* Mention #4482 in new features

Co-authored-by: Michael Osthege <michael.osthege@outlook.com>
---
 RELEASE-NOTES.md         |  1 +
 pymc3/math.py            | 12 +++++++++---
 pymc3/tests/test_math.py | 19 ++++++++++++++++++-
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index f833e6adf76..15dd545a195 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -5,6 +5,7 @@
 + ...
 
 ### New Features
++ `pm.math.cartesian` can now handle inputs that are themselves >1D (see [#4482](https://github.com/pymc-devs/pymc3/pull/4482)).
 + ...
 
 ### Maintenance
diff --git a/pymc3/math.py b/pymc3/math.py
index aff54d13b71..fc176b4185d 100644
--- a/pymc3/math.py
+++ b/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):
diff --git a/pymc3/tests/test_math.py b/pymc3/tests/test_math.py
index b31319021fd..5e73fb9c0f3 100644
--- a/pymc3/tests/test_math.py
+++ b/pymc3/tests/test_math.py
@@ -71,7 +71,24 @@ def test_cartesian():
         ]
     )
     auto_cart = cartesian(a, b, c)
-    np.testing.assert_array_almost_equal(manual_cartesian, auto_cart)
+    np.testing.assert_array_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_equal(manual_cartesian, auto_cart)
 
 
 def test_kron_dot():