diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ddd33b3008..d55b09e4f92 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -97,6 +97,7 @@ This release is compatible with NumPy 2.5. * Fixed a per-call `sycl::queue` leak in `usm_ndarray::get_queue()`/`get_device()` [#3042](https://github.com/IntelPython/dpnp/pull/3042) * Fixed `dpnp.linspace` returning `nan` for equal infinite endpoints [#3043](https://github.com/IntelPython/dpnp/pull/3043) * Fixed `dpnp.cumsum`, `dpnp.cumprod`, and their `nan`/`cumulative_*` variants (including `dpnp.tensor.cumulative_sum`/`cumulative_prod`) silently returning incorrect results when accumulating along an axis of an array with more than one row [#3063](https://github.com/IntelPython/dpnp/pull/3063) +* Fixed `dpnp.einsum` returning a result whose memory layout differs from NumPy for the default `order="K"`, and ignoring `out` and `order` for a contraction over a size-0 dimension [#3058](https://github.com/IntelPython/dpnp/pull/3058) ### Security diff --git a/dpnp/dpnp_utils/dpnp_utils_einsum.py b/dpnp/dpnp_utils/dpnp_utils_einsum.py index 842596bdc32..efa669126a1 100644 --- a/dpnp/dpnp_utils/dpnp_utils_einsum.py +++ b/dpnp/dpnp_utils/dpnp_utils_einsum.py @@ -1039,8 +1039,22 @@ def dpnp_einsum( ) arrays.append(operands[id]) result_dtype = dpnp.result_type(*arrays) if dtype is None else dtype - if order is not None and order in "aA": - order = "F" if all(arr.flags.fnc for arr in arrays) else "C" + # validated here because the view path below skips `dpnp.asarray` + if order is None: + order = "K" + elif not isinstance(order, str): + raise TypeError(f"order must be str, not {type(order).__name__}") + elif len(order) == 1 and order in "afkcAFKC": + order = order.upper() + else: + raise ValueError( + f"order must be one of 'C', 'F', 'A', or 'K' (got '{order}')" + ) + all_f_contiguous = all(arr.flags.f_contiguous for arr in arrays) + if order == "A": + # NumPy uses f_contiguous here, not fnc; they differ for an array that + # is both C- and F-contiguous, such as a 1-D or size-1 one + order = "F" if all_f_contiguous else "C" input_subscripts = [ _parse_ellipsis_subscript(sub, idx, ndim=arr.ndim) @@ -1110,12 +1124,15 @@ def dpnp_einsum( # no more raises if len(operands) >= 2: if any(arr.size == 0 for arr in operands): - return dpnp.zeros( + # NumPy falls back to "C" for "K" here + arr_out = dpnp.zeros( tuple(dimension_dict[label] for label in output_subscript), dtype=result_dtype, + order="C" if order == "K" else order, usm_type=res_usm_type, sycl_queue=exec_q, ) + return dpnp.get_result_array(arr_out, out, casting=casting) # Don't squeeze if unary, because this affects later (in trivial sum) # whether the return is a writeable view. @@ -1226,6 +1243,12 @@ def dpnp_einsum( [dimension_dict[label] for label in output_subscript] ) - arr_out = dpnp.asarray(arr_out, order=order) + # a view is returned for any `order`, the same way NumPy does + if not returns_view: + if order == "K" and optimize is False and not all_f_contiguous: + # only the unoptimized path of NumPy copies into a c-contiguous + # array, the optimized one is matmul-based, as dpnp always is + order = "C" + arr_out = dpnp.asarray(arr_out, order=order) assert returns_view or arr_out.dtype == result_dtype return dpnp.get_result_array(arr_out, out, casting=casting) diff --git a/dpnp/tests/test_linalg.py b/dpnp/tests/test_linalg.py index 9cd9a1b9b8f..fe00f421142 100644 --- a/dpnp/tests/test_linalg.py +++ b/dpnp/tests/test_linalg.py @@ -1690,6 +1690,163 @@ def test_path(self): assert expected[0] == result[0] assert expected[1] == result[1] + @pytest.mark.parametrize( + "subscripts, shape1, shape2", + [ + ("lkz,lxpq->kxpqz", (3, 2, 2), (3, 1, 6, 6)), + ("lkz,lxpq->kxpqz", (4, 3, 2), (4, 2, 5, 5)), + ("ij,jk->ik", (4, 5), (5, 6)), + ("ijk,ikl->ijl", (2, 3, 4), (2, 4, 5)), + ("lk,lpq->kpq", (3, 2), (3, 6, 6)), + ], + ) + def test_contraction_order_k(self, subscripts, shape1, shape2): + # for order="K" (the default), a contraction is materialized into a + # newly allocated array, so the result is c-contiguous when the + # operands are, matching NumPy + a = generate_random_numpy_array(shape1) + b = generate_random_numpy_array(shape2) + ia, ib = dpnp.array(a), dpnp.array(b) + + result = dpnp.einsum(subscripts, ia, ib) + expected = numpy.einsum(subscripts, a, b) + assert result.flags.c_contiguous == expected.flags.c_contiguous + assert result.flags.c_contiguous + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("optimize", [False, True, "greedy", "optimal"]) + @pytest.mark.parametrize("order", ["C", "F", "A", "K", None]) + @pytest.mark.parametrize("order1", ["C", "F"]) + @pytest.mark.parametrize("order2", ["C", "F"]) + def test_contraction_order(self, optimize, order, order1, order2): + if order is None and optimize is not False: + pytest.skip("numpy raises AttributeError for order=None here") + a = generate_random_numpy_array((4, 5), order=order1) + b = generate_random_numpy_array((5, 6), order=order2) + ia, ib = dpnp.array(a), dpnp.array(b) + + result = dpnp.einsum( + "ij,jk->ik", ia, ib, order=order, optimize=optimize + ) + expected = numpy.einsum( + "ij,jk->ik", a, b, order=order, optimize=optimize + ) + assert result.flags.c_contiguous == expected.flags.c_contiguous + assert result.flags.f_contiguous == expected.flags.f_contiguous + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("optimize", [True, "greedy", "optimal"]) + @pytest.mark.parametrize( + "subscripts, shapes", + [ + ("...ft,mf->...mt", [(2, 3, 5), (4, 3)]), + ("lk,lpq->kpq", [(3, 4), (3, 5, 6)]), + ("ijk,jkl->il", [(2, 3, 4), (3, 4, 5)]), + ("ij,jk,kl->il", [(4, 5), (5, 6), (6, 7)]), + ], + ) + @pytest.mark.parametrize("order1", ["C", "F"]) + @pytest.mark.parametrize("order2", ["C", "F"]) + def test_contraction_order_k_optimize( + self, optimize, subscripts, shapes, order1, order2 + ): + # NumPy only copies the result into a c-contiguous array on its + # unoptimized path; an optimized one is matmul-based, as dpnp always + # is, and keeps the permuted layout that the contraction produces + orders = [order1, order2] + ["C"] * (len(shapes) - 2) + arrays = [ + generate_random_numpy_array(shape, order=o) + for shape, o in zip(shapes, orders) + ] + iarrays = [dpnp.array(a, order=o) for a, o in zip(arrays, orders)] + + result = dpnp.einsum(subscripts, *iarrays, optimize=optimize) + expected = numpy.einsum(subscripts, *arrays, optimize=optimize) + assert result.flags.c_contiguous == expected.flags.c_contiguous + assert result.flags.f_contiguous == expected.flags.f_contiguous + assert_dtype_allclose(result, expected) + + def test_contraction_order_a_trivial(self): + # an operand that is both c- and f-contiguous (here 1-D) is + # f_contiguous, so order="A" resolves to "F" as it does in NumPy + a = generate_random_numpy_array(4) + b = generate_random_numpy_array((4, 5, 6), order="F") + ia, ib = dpnp.array(a), dpnp.array(b, order="F") + + result = dpnp.einsum("i,ijk->jk", ia, ib, order="A") + expected = numpy.einsum("i,ijk->jk", a, b, order="A") + assert result.flags.c_contiguous == expected.flags.c_contiguous + assert result.flags.f_contiguous == expected.flags.f_contiguous + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("order", ["C", "F", "A", "K", None]) + def test_empty_operand_order(self, order): + # a contraction over a size-0 dimension is all zeros, and `order` is + # honored for it as it is for a non-empty one + a = numpy.ones((2, 0)) + b = numpy.ones((0, 4)) + ia, ib = dpnp.array(a), dpnp.array(b) + + result = dpnp.einsum("ij,jk->ik", ia, ib, order=order) + expected = numpy.einsum("ij,jk->ik", a, b, order=order) + assert result.flags.c_contiguous == expected.flags.c_contiguous + assert result.flags.f_contiguous == expected.flags.f_contiguous + assert_dtype_allclose(result, expected) + + def test_empty_operand_out(self): + # `out` is filled with zeros and returned for a size-0 contraction + a = numpy.ones((2, 0)) + b = numpy.ones((0, 4)) + ia, ib = dpnp.array(a), dpnp.array(b) + iout = dpnp.full((2, 4), 9.0) + out = numpy.full((2, 4), 9.0) + + result = dpnp.einsum("ij,jk->ik", ia, ib, out=iout) + expected = numpy.einsum("ij,jk->ik", a, b, out=out) + assert result is iout + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("subscripts", ["ij->ji", "ij->ij", "ii->i"]) + @pytest.mark.parametrize("order", ["C", "F", "A", "K", None]) + def test_unary_view_order(self, subscripts, order): + # a single-operand einsum with no summed index returns a view of the + # operand for every value of `order`, as it does in NumPy + # the dtype is pinned because the strides below scale with itemsize + a = generate_random_numpy_array((4, 4), dtype=dpnp.default_float_type()) + ia = dpnp.array(a) + + result = dpnp.einsum(subscripts, ia, order=order) + expected = numpy.einsum(subscripts, a, order=order) + assert result.get_array()._pointer == ia.get_array()._pointer + assert result.strides == expected.strides + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("subscripts", ["ij->ji", "ii->i"]) + @pytest.mark.parametrize("order", ["C", "F", "A", "K"]) + def test_unary_view_is_writeable(self, subscripts, order): + # the view returned for a unary einsum without summation is writeable, + # so an assignment through it is visible in the operand + a = generate_random_numpy_array((4, 4)) + ia = dpnp.array(a) + + result = dpnp.einsum(subscripts, ia, order=order) + result[...] = 0 + expected = numpy.einsum(subscripts, a, order=order) + expected[...] = 0 + assert_dtype_allclose(ia, a) + + @pytest.mark.parametrize("order", ["W", "w", "", "CF"]) + def test_order_error(self, order): + a = dpnp.ones((3, 3)) + # a unary einsum without summation returns a view without going + # through dpnp.asarray, so `order` is validated up front + assert_raises(ValueError, dpnp.einsum, "ii->i", a, order=order) + assert_raises(ValueError, dpnp.einsum, "ij,jk->ik", a, a, order=order) + + def test_order_type_error(self): + a = dpnp.ones((3, 3)) + assert_raises(TypeError, dpnp.einsum, "ii->i", a, order=1) + class TestInv: @pytest.mark.parametrize(