diff --git a/CHANGELOG.md b/CHANGELOG.md index 09961c7ccd3..3f0d3481821 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -95,6 +95,7 @@ This release is compatible with NumPy 2.5. * Fixed `dpnp.insert` silently ignoring out-of-bounds negative indices in a multi-element `obj`, so a mix of in-bounds and out-of-bounds indices now consistently raises `IndexError` [#3041](https://github.com/IntelPython/dpnp/pull/3041) * 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.einsum` returns a result whose memory layout differs from NumPy for the default `order="K"`, and ignores `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..40873628175 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,16 @@ def dpnp_einsum( # no more raises if len(operands) >= 2: if any(arr.size == 0 for arr in operands): - return dpnp.zeros( + # every term of the sum is empty, so the result is all zeros; + # "K" has no layout to keep here, and NumPy falls back to "C" + 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 +1244,14 @@ def dpnp_einsum( [dimension_dict[label] for label in output_subscript] ) - arr_out = dpnp.asarray(arr_out, order=order) + # a unary einsum without summation returns a view, as NumPy does for every + # `order` + if not returns_view: + if order == "K" and not all_f_contiguous: + # NumPy copies the result into a new c-contiguous array, while + # the matmul above leaves a permuted one; for all-f-contiguous + # operands it keeps a layout chosen per contraction, so "K" stays + 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..dd2c6e6dcbd 100644 --- a/dpnp/tests/test_linalg.py +++ b/dpnp/tests/test_linalg.py @@ -1690,6 +1690,125 @@ 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("order", ["C", "F", "A", "K", None]) + @pytest.mark.parametrize("order1", ["C", "F"]) + @pytest.mark.parametrize("order2", ["C", "F"]) + def test_contraction_order(self, order, order1, order2): + 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) + 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_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(