Merge pull request #21262 from ganesh-k13/kron_21257_ma

ENH: Masked Array support for `np.kron`

2 files changed, 23 insertions, 17 deletions

diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
index 581da0598..8371747c0 100644
--- a/numpy/lib/shape_base.py
+++ b/numpy/lib/shape_base.py
@@ -1159,9 +1159,11 @@ def kron(a, b):
     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
+    a_arr = a.reshape(a.size, 1)
+    b_arr = b.reshape(1, b.size)
+    is_any_mat = isinstance(a_arr, matrix) or isinstance(b_arr, matrix)
+    # In case of `mat`, convert result to `array`
+    result = _nx.multiply(a_arr, b_arr, subok=(not is_any_mat))
 
     # Reshape back
     result = result.reshape(as_+bs)
@@ -1169,13 +1171,7 @@ def kron(a, b):
     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)
-    wrapper = get_array_wrap(a, b)
-    if wrapper is not None:
-        result = wrapper(result)
-    return result
+    return result if not is_any_mat else matrix(result, copy=False)
 
 
 def _tile_dispatcher(A, reps):
diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index 564cdfeea..07960627b 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -646,21 +646,31 @@ class TestSqueeze:
 class TestKron:
     def test_return_type(self):
         class myarray(np.ndarray):
-            __array_priority__ = 0.0
+            __array_priority__ = 1.0
 
         a = np.ones([2, 2])
         ma = myarray(a.shape, a.dtype, a.data)
         assert_equal(type(kron(a, a)), np.ndarray)
         assert_equal(type(kron(ma, ma)), myarray)
-        assert_equal(type(kron(a, ma)), np.ndarray)
+        assert_equal(type(kron(a, ma)), myarray)
         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])
+    @pytest.mark.parametrize(
+        "array_class", [np.asarray, np.mat]
+    )
+    def test_kron_smoke(self, array_class):
+        a = array_class(np.ones([3, 3]))
+        b = array_class(np.ones([3, 3]))
+        k = array_class(np.ones([9, 9]))
+
+        assert_array_equal(np.kron(a, b), k)
+
+    def test_kron_ma(self):
+        x = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]])
+        k = np.ma.array(np.diag([1, 4, 4, 16]),
+                mask=~np.array(np.identity(4), dtype=bool))
 
-        assert np.array_equal(np.kron(a, b), k), "Smoke test for kron failed"
+        assert_array_equal(k, np.kron(x, x))
 
     @pytest.mark.parametrize(
         "shape_a,shape_b", [