Skip to content

Commit 1dfef7f

Browse files
committed
Fix landed
1 parent 5da183e commit 1dfef7f

3 files changed

Lines changed: 117 additions & 3 deletions

File tree

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ This release is compatible with NumPy 2.5.
9494
* 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)
9595
* 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)
9696
* Fixed a per-call `sycl::queue` leak in `usm_ndarray::get_queue()`/`get_device()` [#3042](https://github.com/IntelPython/dpnp/pull/3042)
97+
* 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)
9798

9899
### Security
99100

dpnp/dpnp_utils/dpnp_utils_einsum.py

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1039,8 +1039,21 @@ def dpnp_einsum(
10391039
)
10401040
arrays.append(operands[id])
10411041
result_dtype = dpnp.result_type(*arrays) if dtype is None else dtype
1042-
if order is not None and order in "aA":
1043-
order = "F" if all(arr.flags.fnc for arr in arrays) else "C"
1042+
# validated here because the view path below skips `dpnp.asarray`
1043+
if order is None:
1044+
order = "K"
1045+
elif not isinstance(order, str):
1046+
raise TypeError(f"order must be str, not {type(order).__name__}")
1047+
elif len(order) == 1 and order in "afkcAFKC":
1048+
order = order.upper()
1049+
else:
1050+
raise ValueError(
1051+
f"order must be one of 'C', 'F', 'A', or 'K' (got '{order}')"
1052+
)
1053+
if order == "A":
1054+
# NumPy uses f_contiguous here, not fnc; they differ for an array that
1055+
# is both C- and F-contiguous, such as a 1-D or size-1 one
1056+
order = "F" if all(arr.flags.f_contiguous for arr in arrays) else "C"
10441057

10451058
input_subscripts = [
10461059
_parse_ellipsis_subscript(sub, idx, ndim=arr.ndim)
@@ -1226,6 +1239,13 @@ def dpnp_einsum(
12261239
[dimension_dict[label] for label in output_subscript]
12271240
)
12281241

1229-
arr_out = dpnp.asarray(arr_out, order=order)
1242+
# a unary einsum without summation returns a view, as NumPy does for every
1243+
# `order`
1244+
if not returns_view:
1245+
if order == "K" and all(arr.flags.c_contiguous for arr in arrays):
1246+
# NumPy copies the result into a new c-contiguous array, while
1247+
# the matmul above leaves a permuted one; other layouts vary
1248+
order = "C"
1249+
arr_out = dpnp.asarray(arr_out, order=order)
12301250
assert returns_view or arr_out.dtype == result_dtype
12311251
return dpnp.get_result_array(arr_out, out, casting=casting)

dpnp/tests/test_linalg.py

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1690,6 +1690,99 @@ def test_path(self):
16901690
assert expected[0] == result[0]
16911691
assert expected[1] == result[1]
16921692

1693+
@pytest.mark.parametrize(
1694+
"subscripts, shape1, shape2",
1695+
[
1696+
("lkz,lxpq->kxpqz", (3, 2, 2), (3, 1, 6, 6)),
1697+
("lkz,lxpq->kxpqz", (4, 3, 2), (4, 2, 5, 5)),
1698+
("ij,jk->ik", (4, 5), (5, 6)),
1699+
("ijk,ikl->ijl", (2, 3, 4), (2, 4, 5)),
1700+
("lk,lpq->kpq", (3, 2), (3, 6, 6)),
1701+
],
1702+
)
1703+
def test_contraction_order_k(self, subscripts, shape1, shape2):
1704+
# for order="K" (the default), a contraction is materialized into a
1705+
# newly allocated array, so the result is c-contiguous when the
1706+
# operands are, matching NumPy
1707+
a = generate_random_numpy_array(shape1)
1708+
b = generate_random_numpy_array(shape2)
1709+
ia, ib = dpnp.array(a), dpnp.array(b)
1710+
1711+
result = dpnp.einsum(subscripts, ia, ib)
1712+
expected = numpy.einsum(subscripts, a, b)
1713+
assert result.flags.c_contiguous == expected.flags.c_contiguous
1714+
assert result.flags.c_contiguous
1715+
assert_dtype_allclose(result, expected)
1716+
1717+
@pytest.mark.parametrize("order", ["C", "F", "A"])
1718+
@pytest.mark.parametrize("order1", ["C", "F"])
1719+
@pytest.mark.parametrize("order2", ["C", "F"])
1720+
def test_contraction_order(self, order, order1, order2):
1721+
# order="K" is covered by test_contraction_order_k; for operands that
1722+
# are not all c-contiguous NumPy keeps a layout chosen per contraction
1723+
a = generate_random_numpy_array((4, 5), order=order1)
1724+
b = generate_random_numpy_array((5, 6), order=order2)
1725+
ia, ib = dpnp.array(a), dpnp.array(b)
1726+
1727+
result = dpnp.einsum("ij,jk->ik", ia, ib, order=order)
1728+
expected = numpy.einsum("ij,jk->ik", a, b, order=order)
1729+
assert result.flags.c_contiguous == expected.flags.c_contiguous
1730+
assert result.flags.f_contiguous == expected.flags.f_contiguous
1731+
assert_dtype_allclose(result, expected)
1732+
1733+
def test_contraction_order_a_trivial(self):
1734+
# an operand that is both c- and f-contiguous (here 1-D) is
1735+
# f_contiguous, so order="A" resolves to "F" as it does in NumPy
1736+
a = generate_random_numpy_array(4)
1737+
b = generate_random_numpy_array((4, 5, 6), order="F")
1738+
ia, ib = dpnp.array(a), dpnp.array(b, order="F")
1739+
1740+
result = dpnp.einsum("i,ijk->jk", ia, ib, order="A")
1741+
expected = numpy.einsum("i,ijk->jk", a, b, order="A")
1742+
assert result.flags.c_contiguous == expected.flags.c_contiguous
1743+
assert result.flags.f_contiguous == expected.flags.f_contiguous
1744+
assert_dtype_allclose(result, expected)
1745+
1746+
@pytest.mark.parametrize("subscripts", ["ij->ji", "ij->ij", "ii->i"])
1747+
@pytest.mark.parametrize("order", ["C", "F", "A", "K", None])
1748+
def test_unary_view_order(self, subscripts, order):
1749+
# a single-operand einsum with no summed index returns a view of the
1750+
# operand for every value of `order`, as it does in NumPy
1751+
a = generate_random_numpy_array((4, 4))
1752+
ia = dpnp.array(a)
1753+
1754+
result = dpnp.einsum(subscripts, ia, order=order)
1755+
expected = numpy.einsum(subscripts, a, order=order)
1756+
assert result.get_array()._pointer == ia.get_array()._pointer
1757+
assert result.strides == expected.strides
1758+
assert_dtype_allclose(result, expected)
1759+
1760+
@pytest.mark.parametrize("subscripts", ["ij->ji", "ii->i"])
1761+
@pytest.mark.parametrize("order", ["C", "F", "A", "K"])
1762+
def test_unary_view_is_writeable(self, subscripts, order):
1763+
# the view returned for a unary einsum without summation is writeable,
1764+
# so an assignment through it is visible in the operand
1765+
a = generate_random_numpy_array((4, 4))
1766+
ia = dpnp.array(a)
1767+
1768+
result = dpnp.einsum(subscripts, ia, order=order)
1769+
result[...] = 0
1770+
expected = numpy.einsum(subscripts, a, order=order)
1771+
expected[...] = 0
1772+
assert_dtype_allclose(ia, a)
1773+
1774+
@pytest.mark.parametrize("order", ["W", "w", "", "CF"])
1775+
def test_order_error(self, order):
1776+
a = dpnp.ones((3, 3))
1777+
# a unary einsum without summation returns a view without going
1778+
# through dpnp.asarray, so `order` is validated up front
1779+
assert_raises(ValueError, dpnp.einsum, "ii->i", a, order=order)
1780+
assert_raises(ValueError, dpnp.einsum, "ij,jk->ik", a, a, order=order)
1781+
1782+
def test_order_type_error(self):
1783+
a = dpnp.ones((3, 3))
1784+
assert_raises(TypeError, dpnp.einsum, "ii->i", a, order=1)
1785+
16931786

16941787
class TestInv:
16951788
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)