Merge pull request #18708 from DCtheTall/meshgrid-test

TST: add tests for using np.meshgrid for higher dimensional grids.

1 files changed, 21 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index afcb81eff..761ea83a3 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -2307,6 +2307,27 @@ class TestMeshgrid:
         assert_equal(x[0, :], 0)
         assert_equal(x[1, :], X)
 
+    def test_nd_shape(self):
+        a, b, c, d, e = np.meshgrid(*([0] * i for i in range(1, 6)))
+        expected_shape = (2, 1, 3, 4, 5)
+        assert_equal(a.shape, expected_shape)
+        assert_equal(b.shape, expected_shape)
+        assert_equal(c.shape, expected_shape)
+        assert_equal(d.shape, expected_shape)
+        assert_equal(e.shape, expected_shape)
+
+    def test_nd_values(self):
+        a, b, c = np.meshgrid([0], [1, 2], [3, 4, 5])
+        assert_equal(a, [[[0, 0, 0]], [[0, 0, 0]]])
+        assert_equal(b, [[[1, 1, 1]], [[2, 2, 2]]])
+        assert_equal(c, [[[3, 4, 5]], [[3, 4, 5]]])
+
+    def test_nd_indexing(self):
+        a, b, c = np.meshgrid([0], [1, 2], [3, 4, 5], indexing='ij')
+        assert_equal(a, [[[0, 0, 0], [0, 0, 0]]])
+        assert_equal(b, [[[1, 1, 1], [2, 2, 2]]])
+        assert_equal(c, [[[3, 4, 5], [3, 4, 5]]])
+
 
 class TestPiecewise: