From 664f853c8dc49ddc4859b43e47c73334fba5aca8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Gauthier-Clerc?= Date: Thu, 6 Nov 2025 20:02:06 +0100 Subject: [PATCH 1/2] fix bad delegation behaviour with atleast_3d and improve atleast_nd unittests. --- src/array_api_extra/_delegation.py | 2 +- tests/test_funcs.py | 189 ++++++++++++++++++++--------- 2 files changed, 132 insertions(+), 59 deletions(-) diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 289d21e4..4cdf255a 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -67,7 +67,7 @@ def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array if xp is None: xp = array_namespace(x) - if 1 <= ndim <= 3 and ( + if 1 <= ndim <= 2 and ( is_numpy_namespace(xp) or is_jax_namespace(xp) or is_dask_namespace(xp) diff --git a/tests/test_funcs.py b/tests/test_funcs.py index ff050468..9e6b7296 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -54,6 +54,8 @@ lazy_xp_function(setdiff1d, jax_jit=False) lazy_xp_function(sinc) +NestedFloatList = list[float] | list["NestedFloatList"] + class TestApplyWhere: @staticmethod @@ -291,68 +293,139 @@ def test_0D(self, xp: ModuleType): y = atleast_nd(x, ndim=5) xp_assert_equal(y, xp.ones((1, 1, 1, 1, 1))) - def test_1D(self, xp: ModuleType): - x = xp.asarray([0, 1]) - - y = atleast_nd(x, ndim=0) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=1) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - xp_assert_equal(y, xp.asarray([[0, 1]])) - - y = atleast_nd(x, ndim=5) - xp_assert_equal(y, xp.asarray([[[[[0, 1]]]]])) - - def test_2D(self, xp: ModuleType): - x = xp.asarray([[3.0]]) - - y = atleast_nd(x, ndim=0) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=3) - xp_assert_equal(y, 3 * xp.ones((1, 1, 1))) - - y = atleast_nd(x, ndim=5) - xp_assert_equal(y, 3 * xp.ones((1, 1, 1, 1, 1))) - - def test_3D(self, xp: ModuleType): - x = xp.asarray([[[3.0], [2.0]]]) - - y = atleast_nd(x, ndim=0) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=3) - xp_assert_equal(y, x) - - y = atleast_nd(x, ndim=5) - xp_assert_equal(y, xp.asarray([[[[[3.0], [2.0]]]]])) - - def test_5D(self, xp: ModuleType): - x = xp.ones((1, 1, 1, 1, 1)) - - y = atleast_nd(x, ndim=0) - xp_assert_equal(y, x) + @pytest.mark.parametrize( + ("x_data", "ndim", "expected_data"), + [ + # --- size-1 vector --- + ([3.0], 0, [3.0]), + ([3.0], 1, [3.0]), + ([3.0], 2, [[3.0]]), + ([3.0], 3, [[[3.0]]]), + ([3.0], 5, [[[[[3.0]]]]]), + # --- size-2 vector --- + ([0.0, 1.0], 0, [0.0, 1.0]), + ([0.0, 1.0], 1, [0.0, 1.0]), + ([0.0, 1.0], 2, [[0.0, 1.0]]), + ([0.0, 1.0], 5, [[[[[0.0, 1.0]]]]]), + ], + ) + def test_1D( + self, + x_data: NestedFloatList, + ndim: int, + expected_data: NestedFloatList, + xp: ModuleType, + ): + x = xp.asarray(x_data) + expected = xp.asarray(expected_data) + y = atleast_nd(x, ndim=ndim) + xp_assert_equal(y, expected) - y = atleast_nd(x, ndim=4) - xp_assert_equal(y, x) + @pytest.mark.parametrize( + ("x_data", "ndim", "expected_data"), + [ + # --- size-1 vector --- + ([[3.0]], 0, [[3.0]]), + ([[3.0]], 1, [[3.0]]), + ([[3.0]], 2, [[3.0]]), + ([[3.0]], 3, [[[3.0]]]), + ([[3.0]], 5, [[[[[3.0]]]]]), + # --- size-2 vector --- + ([[0.0], [1.0]], 0, [[0.0], [1.0]]), + ([[0.0, 1.0]], 1, [[0.0, 1.0]]), + ([[0.0, 1.0]], 2, [[0.0, 1.0]]), + ([[0.0], [1.0]], 3, [[[0.0], [1.0]]]), + ([[0.0, 1.0]], 5, [[[[[0.0, 1.0]]]]]), + ], + ) + def test_2D( + self, + x_data: NestedFloatList, + ndim: int, + expected_data: NestedFloatList, + xp: ModuleType, + ): + x = xp.asarray(x_data) + expected = xp.asarray(expected_data) + y = atleast_nd(x, ndim=ndim) + xp_assert_equal(y, expected) - y = atleast_nd(x, ndim=5) - xp_assert_equal(y, x) + @pytest.mark.parametrize( + ("x_data", "ndim", "expected_data"), + [ + ([[[0.0]], [[1.0]]], 0, [[[0.0]], [[1.0]]]), + ([[[0.0], [1.0]]], 1, [[[0.0], [1.0]]]), + ([[[0.0, 1.0]]], 2, [[[0.0, 1.0]]]), + ([[[0.0]], [[1.0]]], 3, [[[0.0]], [[1.0]]]), + ([[[0.0], [1.0]]], 5, [[[[[0.0], [1.0]]]]]), + ], + ) + def test_3D( + self, + x_data: NestedFloatList, + ndim: int, + expected_data: NestedFloatList, + xp: ModuleType, + ): + x = xp.asarray(x_data) + expected = xp.asarray(expected_data) + y = atleast_nd(x, ndim=ndim) + xp_assert_equal(y, expected) - y = atleast_nd(x, ndim=6) - xp_assert_equal(y, xp.ones((1, 1, 1, 1, 1, 1))) + @pytest.mark.parametrize( + ("x_data", "ndim", "expected_data"), + [ + ([[[[3.0], [2.0]]]], 0, [[[[3.0], [2.0]]]]), + ([[[[3.0, 2.0]]]], 2, [[[[3.0, 2.0]]]]), + ([[[[3.0]], [[2.0]]]], 4, [[[[3.0]], [[2.0]]]]), + ([[[[3.0]]], [[[2.0]]]], 5, [[[[[3.0]]], [[[2.0]]]]]), + ], + ) + def test_4D( + self, + x_data: NestedFloatList, + ndim: int, + expected_data: NestedFloatList, + xp: ModuleType, + ): + x = xp.asarray(x_data) + expected = xp.asarray(expected_data) + y = atleast_nd(x, ndim=ndim) + xp_assert_equal(y, expected) - y = atleast_nd(x, ndim=9) - xp_assert_equal(y, xp.ones((1, 1, 1, 1, 1, 1, 1, 1, 1))) + @pytest.mark.parametrize( + ("x_data", "ndim", "expected_data"), + [ + ([[[[[3.0]], [[2.0]], [[1.0]]]]], 0, [[[[[3.0]], [[2.0]], [[1.0]]]]]), + ([[[[[3.0, 2.0, 6.0]]]]], 2, [[[[[3.0, 2.0, 6.0]]]]]), + ( + [[[[[3.0]]], [[[2.0]]], [[[1.0]]]]], + 4, + [[[[[3.0]]], [[[2.0]]], [[[1.0]]]]], + ), + ( + [[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]], + 6, + [[[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]]], + ), + ( + [[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]], + 9, + [[[[[[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]]]]]], + ), + ], + ) + def test_5D( + self, + x_data: NestedFloatList, + ndim: int, + expected_data: NestedFloatList, + xp: ModuleType, + ): + x = xp.asarray(x_data) + expected = xp.asarray(expected_data) + y = atleast_nd(x, ndim=ndim) + xp_assert_equal(y, expected) def test_device(self, xp: ModuleType, device: Device): x = xp.asarray([1, 2, 3], device=device) From 6ba1c518e07c692c06e3d0612a398c4d05da0c6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Gauthier-Clerc?= Date: Sun, 9 Nov 2025 17:42:16 +0100 Subject: [PATCH 2/2] reduce atleast_test verbosity. --- tests/test_funcs.py | 216 ++++++++++++++++++++++++-------------------- 1 file changed, 116 insertions(+), 100 deletions(-) diff --git a/tests/test_funcs.py b/tests/test_funcs.py index 9e6b7296..6631cc32 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -294,147 +294,163 @@ def test_0D(self, xp: ModuleType): xp_assert_equal(y, xp.ones((1, 1, 1, 1, 1))) @pytest.mark.parametrize( - ("x_data", "ndim", "expected_data"), + ("input_shape", "ndim", "expected_shape"), [ - # --- size-1 vector --- - ([3.0], 0, [3.0]), - ([3.0], 1, [3.0]), - ([3.0], 2, [[3.0]]), - ([3.0], 3, [[[3.0]]]), - ([3.0], 5, [[[[[3.0]]]]]), - # --- size-2 vector --- - ([0.0, 1.0], 0, [0.0, 1.0]), - ([0.0, 1.0], 1, [0.0, 1.0]), - ([0.0, 1.0], 2, [[0.0, 1.0]]), - ([0.0, 1.0], 5, [[[[[0.0, 1.0]]]]]), + ((1,), 0, (1,)), + ((5,), 1, (5,)), + ((2,), 2, (1, 2)), + ((3,), 3, (1, 1, 3)), + ((2,), 5, (1, 1, 1, 1, 2)), ], ) - def test_1D( + def test_1D_shapes( self, - x_data: NestedFloatList, + input_shape: tuple[int], ndim: int, - expected_data: NestedFloatList, + expected_shape: tuple[int], xp: ModuleType, ): - x = xp.asarray(x_data) - expected = xp.asarray(expected_data) + n = math.prod(input_shape) + x = xp.reshape(xp.asarray(list(range(n))), input_shape) y = atleast_nd(x, ndim=ndim) - xp_assert_equal(y, expected) - @pytest.mark.parametrize( - ("x_data", "ndim", "expected_data"), - [ - # --- size-1 vector --- - ([[3.0]], 0, [[3.0]]), - ([[3.0]], 1, [[3.0]]), - ([[3.0]], 2, [[3.0]]), - ([[3.0]], 3, [[[3.0]]]), - ([[3.0]], 5, [[[[[3.0]]]]]), - # --- size-2 vector --- - ([[0.0], [1.0]], 0, [[0.0], [1.0]]), - ([[0.0, 1.0]], 1, [[0.0, 1.0]]), - ([[0.0, 1.0]], 2, [[0.0, 1.0]]), - ([[0.0], [1.0]], 3, [[[0.0], [1.0]]]), - ([[0.0, 1.0]], 5, [[[[[0.0, 1.0]]]]]), - ], - ) - def test_2D( - self, - x_data: NestedFloatList, - ndim: int, - expected_data: NestedFloatList, - xp: ModuleType, - ): - x = xp.asarray(x_data) - expected = xp.asarray(expected_data) - y = atleast_nd(x, ndim=ndim) - xp_assert_equal(y, expected) + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_1D_values(self, xp: ModuleType): + x = xp.asarray([0, 1]) + + y = atleast_nd(x, ndim=0) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=1) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + xp_assert_equal(y, xp.asarray([[0, 1]])) + + y = atleast_nd(x, ndim=5) + xp_assert_equal(y, xp.asarray([[[[[0, 1]]]]])) @pytest.mark.parametrize( - ("x_data", "ndim", "expected_data"), + ("input_shape", "ndim", "expected_shape"), [ - ([[[0.0]], [[1.0]]], 0, [[[0.0]], [[1.0]]]), - ([[[0.0], [1.0]]], 1, [[[0.0], [1.0]]]), - ([[[0.0, 1.0]]], 2, [[[0.0, 1.0]]]), - ([[[0.0]], [[1.0]]], 3, [[[0.0]], [[1.0]]]), - ([[[0.0], [1.0]]], 5, [[[[[0.0], [1.0]]]]]), + ((2, 1), 0, (2, 1)), + ((5, 2), 1, (5, 2)), + ((2, 1), 2, (2, 1)), + ((3, 1), 3, (1, 3, 1)), + ((2, 8), 5, (1, 1, 1, 2, 8)), ], ) - def test_3D( + def test_2D_shapes( self, - x_data: NestedFloatList, + input_shape: tuple[int], ndim: int, - expected_data: NestedFloatList, + expected_shape: tuple[int], xp: ModuleType, ): - x = xp.asarray(x_data) - expected = xp.asarray(expected_data) + n = math.prod(input_shape) + x = xp.reshape(xp.asarray(list(range(n))), input_shape) y = atleast_nd(x, ndim=ndim) - xp_assert_equal(y, expected) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_2D_values(self, xp: ModuleType): + x = xp.asarray([[3.0], [4.0]]) + + y = atleast_nd(x, ndim=0) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=3) + xp_assert_equal(y, xp.asarray([[[3.0], [4.0]]])) + + y = atleast_nd(x, ndim=5) + xp_assert_equal(y, xp.asarray([[[[[3.0], [4.0]]]]])) @pytest.mark.parametrize( - ("x_data", "ndim", "expected_data"), + ("input_shape", "ndim", "expected_shape"), [ - ([[[[3.0], [2.0]]]], 0, [[[[3.0], [2.0]]]]), - ([[[[3.0, 2.0]]]], 2, [[[[3.0, 2.0]]]]), - ([[[[3.0]], [[2.0]]]], 4, [[[[3.0]], [[2.0]]]]), - ([[[[3.0]]], [[[2.0]]]], 5, [[[[[3.0]]], [[[2.0]]]]]), + ((2, 1, 1), 0, (2, 1, 1)), + ((1, 5, 2), 1, (1, 5, 2)), + ((2, 1, 1), 2, (2, 1, 1)), + ((1, 3, 1), 3, (1, 3, 1)), + ((2, 8, 1), 5, (1, 1, 2, 8, 1)), ], ) - def test_4D( + def test_3D_shapes( self, - x_data: NestedFloatList, + input_shape: tuple[int], ndim: int, - expected_data: NestedFloatList, + expected_shape: tuple[int], xp: ModuleType, ): - x = xp.asarray(x_data) - expected = xp.asarray(expected_data) + n = math.prod(input_shape) + x = xp.reshape(xp.asarray(list(range(n))), input_shape) y = atleast_nd(x, ndim=ndim) - xp_assert_equal(y, expected) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_3D_values(self, xp: ModuleType): + x = xp.asarray([[[3.0], [2.0]]]) + + y = atleast_nd(x, ndim=0) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=3) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=5) + xp_assert_equal(y, xp.asarray([[[[[3.0], [2.0]]]]])) @pytest.mark.parametrize( - ("x_data", "ndim", "expected_data"), + ("input_shape", "ndim", "expected_shape"), [ - ([[[[[3.0]], [[2.0]], [[1.0]]]]], 0, [[[[[3.0]], [[2.0]], [[1.0]]]]]), - ([[[[[3.0, 2.0, 6.0]]]]], 2, [[[[[3.0, 2.0, 6.0]]]]]), - ( - [[[[[3.0]]], [[[2.0]]], [[[1.0]]]]], - 4, - [[[[[3.0]]], [[[2.0]]], [[[1.0]]]]], - ), - ( - [[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]], - 6, - [[[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]]], - ), - ( - [[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]], - 9, - [[[[[[[[[3.0]], [[1.0]]], [[[2.0]], [[1.0]]], [[[1.0]], [[1.0]]]]]]]]], - ), + ((2, 1, 1, 2, 1), 0, (2, 1, 1, 2, 1)), + ((1, 5, 2, 3, 2), 2, (1, 5, 2, 3, 2)), + ((2, 1, 1, 5, 2), 5, (2, 1, 1, 5, 2)), + ((1, 3, 1, 2, 1), 6, (1, 1, 3, 1, 2, 1)), + ((2, 8, 1, 9, 8), 9, (1, 1, 1, 1, 2, 8, 1, 9, 8)), ], ) - def test_5D( + def test_5D_shapes( self, - x_data: NestedFloatList, + input_shape: tuple[int], ndim: int, - expected_data: NestedFloatList, + expected_shape: tuple[int], xp: ModuleType, ): - x = xp.asarray(x_data) - expected = xp.asarray(expected_data) + n = math.prod(input_shape) + x = xp.reshape(xp.asarray(list(range(n))), input_shape) y = atleast_nd(x, ndim=ndim) - xp_assert_equal(y, expected) - def test_device(self, xp: ModuleType, device: Device): - x = xp.asarray([1, 2, 3], device=device) - assert get_device(atleast_nd(x, ndim=2)) == device + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) - def test_xp(self, xp: ModuleType): - x = xp.asarray(1.0) - y = atleast_nd(x, ndim=1, xp=xp) - xp_assert_equal(y, xp.ones((1,))) + def test_5D_values(self, xp: ModuleType): + x = xp.asarray([[[[[3.0]], [[2.0]]]]]) + + y = atleast_nd(x, ndim=0) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=4) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=5) + xp_assert_equal(y, x) + + y = atleast_nd(x, ndim=6) + xp_assert_equal(y, xp.asarray([[[[[[3.0]], [[2.0]]]]]])) + + y = atleast_nd(x, ndim=9) + xp_assert_equal(y, xp.asarray([[[[[[[[[3.0]], [[2.0]]]]]]]]])) class TestBroadcastShapes: