Fix dpnp.einsum memory-layout

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
 

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)

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(