From 1dfef7fb1d7df8017f445c0ec0223a032053354d Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Wed, 2 Sep 2026 12:14:57 -0500 Subject: [PATCH] Fix landed --- CHANGELOG.md | 1 + dpnp/dpnp_utils/dpnp_utils_einsum.py | 26 +++++++- dpnp/tests/test_linalg.py | 93 ++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e37ae705527..b37eef86a68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -94,6 +94,7 @@ This release is compatible with NumPy 2.5. * Fixed `astype` casting an out-of-range floating point value to a signed narrow integer type saturating to the destination min/max instead of wrapping like NumPy, generalizing the earlier unsigned-only fix [#3033](https://github.com/IntelPython/dpnp/pull/3033) * 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.einsum` returns a non-contiguous result for a contraction over c-contiguous operands with the default `order="K"` [#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..bca5d529c0d 100644 --- a/dpnp/dpnp_utils/dpnp_utils_einsum.py +++ b/dpnp/dpnp_utils/dpnp_utils_einsum.py @@ -1039,8 +1039,21 @@ 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}')" + ) + 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(arr.flags.f_contiguous for arr in arrays) else "C" input_subscripts = [ _parse_ellipsis_subscript(sub, idx, ndim=arr.ndim) @@ -1226,6 +1239,13 @@ 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 all(arr.flags.c_contiguous for arr in arrays): + # NumPy copies the result into a new c-contiguous array, while + # the matmul above leaves a permuted one; other layouts vary + 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..3dd58c17f3b 100644 --- a/dpnp/tests/test_linalg.py +++ b/dpnp/tests/test_linalg.py @@ -1690,6 +1690,99 @@ 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"]) + @pytest.mark.parametrize("order1", ["C", "F"]) + @pytest.mark.parametrize("order2", ["C", "F"]) + def test_contraction_order(self, order, order1, order2): + # order="K" is covered by test_contraction_order_k; for operands that + # are not all c-contiguous NumPy keeps a layout chosen per contraction + 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("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 + a = generate_random_numpy_array((4, 4)) + 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(