Merge pull request #21232 from ganesh-k13/bug_21051_kron

BUG: Fixes `ValueError` in `np.kron`

2 files changed, 46 insertions, 11 deletions

diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
index b600b70f6..581da0598 100644
--- a/numpy/lib/shape_base.py
+++ b/numpy/lib/shape_base.py
@@ -1142,25 +1142,33 @@ def kron(a, b):
     b = asanyarray(b)
     a = array(a, copy=False, subok=True, ndmin=b.ndim)
     ndb, nda = b.ndim, a.ndim
+    nd = max(ndb, nda)
+
     if (nda == 0 or ndb == 0):
         return _nx.multiply(a, b)
+
     as_ = a.shape
     bs = b.shape
     if not a.flags.contiguous:
         a = reshape(a, as_)
     if not b.flags.contiguous:
         b = reshape(b, bs)
-    nd = ndb
-    if (ndb != nda):
-        if (ndb > nda):
-            as_ = (1,)*(ndb-nda) + as_
-        else:
-            bs = (1,)*(nda-ndb) + bs
-            nd = nda
-    result = outer(a, b).reshape(as_+bs)
-    axis = nd-1
-    for _ in range(nd):
-        result = concatenate(result, axis=axis)
+
+    # Equalise the shapes by prepending smaller one with 1s
+    as_ = (1,)*max(0, ndb-nda) + as_
+    bs = (1,)*max(0, nda-ndb) + bs
+
+    # Compute the product
+    a_arr = _nx.asarray(a).reshape(a.size, 1)
+    b_arr = _nx.asarray(b).reshape(1, b.size)
+    result = a_arr * b_arr
+
+    # Reshape back
+    result = result.reshape(as_+bs)
+    transposer = _nx.arange(nd*2).reshape([2, nd]).ravel(order='f')
+    result = result.transpose(transposer)
+    result = result.reshape(_nx.multiply(as_, bs))
+
     wrapper = get_array_prepare(a, b)
     if wrapper is not None:
         result = wrapper(result)
diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index a148e53da..564cdfeea 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -655,6 +655,33 @@ class TestKron:
         assert_equal(type(kron(a, ma)), np.ndarray)
         assert_equal(type(kron(ma, a)), myarray)
 
+    def test_kron_smoke(self):
+        a = np.ones([3, 3])
+        b = np.ones([3, 3])
+        k = np.ones([9, 9])
+
+        assert np.array_equal(np.kron(a, b), k), "Smoke test for kron failed"
+
+    @pytest.mark.parametrize(
+        "shape_a,shape_b", [
+            ((1, 1), (1, 1)),
+            ((1, 2, 3), (4, 5, 6)),
+            ((2, 2), (2, 2, 2)),
+            ((1, 0), (1, 1)),
+            ((2, 0, 2), (2, 2)),
+            ((2, 0, 0, 2), (2, 0, 2)),
+        ])
+    def test_kron_shape(self, shape_a, shape_b):
+        a = np.ones(shape_a)
+        b = np.ones(shape_b)
+        normalised_shape_a = (1,) * max(0, len(shape_b)-len(shape_a)) + shape_a
+        normalised_shape_b = (1,) * max(0, len(shape_a)-len(shape_b)) + shape_b
+        expected_shape = np.multiply(normalised_shape_a, normalised_shape_b)
+
+        k = np.kron(a, b)
+        assert np.array_equal(
+                k.shape, expected_shape), "Unexpected shape from kron"
+
 
 class TestTile:
     def test_basic(self):