diff options
| -rw-r--r-- | benchmarks/benchmarks/bench_shape_base.py | 18 | ||||
| -rw-r--r-- | numpy/lib/shape_base.py | 30 | ||||
| -rw-r--r-- | numpy/lib/tests/test_shape_base.py | 27 |
3 files changed, 64 insertions, 11 deletions
diff --git a/benchmarks/benchmarks/bench_shape_base.py b/benchmarks/benchmarks/bench_shape_base.py index 0c7dc4e72..375c43dcf 100644 --- a/benchmarks/benchmarks/bench_shape_base.py +++ b/benchmarks/benchmarks/bench_shape_base.py @@ -134,3 +134,21 @@ class Block3D(Benchmark): # Retain old benchmark name for backward compat time_3d.benchmark_name = "bench_shape_base.Block.time_3d" + + +class Kron(Benchmark): + """Benchmarks for Kronecker product of two arrays""" + + def setup(self): + self.large_arr = np.random.random((10,) * 4) + self.large_mat = np.mat(np.random.random((100, 100))) + self.scalar = 7 + + def time_arr_kron(self): + np.kron(self.large_arr, self.large_arr) + + def time_scalar_kron(self): + np.kron(self.large_arr, self.scalar) + + def time_mat_kron(self): + np.kron(self.large_mat, self.large_mat) 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): |