diff --git a/dask-xfails.txt b/dask-xfails.txt index fc650923..2976f568 100644 --- a/dask-xfails.txt +++ b/dask-xfails.txt @@ -130,11 +130,6 @@ array_api_tests/test_linalg.py::test_qr array_api_tests/test_manipulation_functions.py::test_roll # 2025.12 support -array_api_tests/test_has_names.py::test_has_names[manipulation-broadcast_shapes] -array_api_tests/test_signatures.py::test_func_signature[broadcast_shapes] -array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_broadcast_shapes -array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_empty -array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_error array_api_tests/test_searching_functions.py::test_searchsorted_with_scalars array_api_tests/test_linalg.py::test_eig diff --git a/src/array_api_compat/dask/array/_aliases.py b/src/array_api_compat/dask/array/_aliases.py index 54d323b2..b30cf665 100644 --- a/src/array_api_compat/dask/array/_aliases.py +++ b/src/array_api_compat/dask/array/_aliases.py @@ -13,6 +13,7 @@ from typing_extensions import TypeIs import dask.array as da +from dask.array.core import broadcast_shapes import numpy as np from numpy import bool_ as bool from numpy import ( @@ -351,6 +352,7 @@ def count_nonzero( __all__ = [ + "broadcast_shapes", "count_nonzero", "bool", "int8", "int16", "int32", "int64", diff --git a/tests/test_all.py b/tests/test_all.py index d9350ce7..59343539 100644 --- a/tests/test_all.py +++ b/tests/test_all.py @@ -231,7 +231,7 @@ XFAILS = { ("numpy", ""): ["from_dlpack"] if np.__version__ < "1.23" else [], - ("dask.array", ""): ["from_dlpack", "take_along_axis", "broadcast_shapes"], + ("dask.array", ""): ["from_dlpack", "take_along_axis"], ("dask.array", "linalg"): [ "cross", "det", diff --git a/tests/test_dask.py b/tests/test_dask.py index 4200e5b7..c446a4a1 100644 --- a/tests/test_dask.py +++ b/tests/test_dask.py @@ -80,6 +80,22 @@ def test_asarray_no_compute(xp): xp.asarray(a, copy=False, dtype=a.dtype) +@pytest.mark.parametrize("unknown", [False, True]) +def test_broadcast_shapes_no_compute(xp, unknown): + a = da.arange(6, chunks=2) + if unknown: + a = a[a > 0] + chunks = a.chunks + with assert_no_compute(): + shape = xp.broadcast_shapes(a.shape, (1,)) + assert isinstance(shape, tuple) + if unknown: + assert np.isnan(shape[0]) + else: + assert shape == (6,) + assert a.chunks == chunks + + @pytest.mark.parametrize("copy", [True, False]) def test_astype_no_compute(xp, copy): with assert_no_compute():