From ce69b4d75f6179d4ee9523785e3aa158d4475f9f Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 12 Mar 2025 20:35:55 +0530 Subject: [PATCH 1/3] Modify atleast_Nd to accept only one positional argument --- pytensor/tensor/basic.py | 30 ++++++++++++------------------ tests/tensor/test_basic.py | 3 ++- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index e30887cfe3..a6055824db 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -4355,28 +4355,22 @@ def empty_like( def atleast_Nd( - *arys: np.ndarray | TensorVariable, n: int = 1, left: bool = True + arry: np.ndarray | TensorVariable, n: int = 1, left: bool = True ) -> TensorVariable: - """Convert inputs to arrays with at least `n` dimensions.""" - res = [] - for ary in arys: - ary = as_tensor(ary) + """Convert input to an array with at least `n` dimensions.""" - if ary.ndim >= n: - result = ary - else: - result = ( - shape_padleft(ary, n - ary.ndim) - if left - else shape_padright(ary, n - ary.ndim) - ) + arry = as_tensor(arry) - res.append(result) - - if len(res) == 1: - return res[0] + if arry.ndim >= n: + result = arry else: - return res + result = ( + shape_padleft(arry, n - arry.ndim) + if left + else shape_padright(arry, n - arry.ndim) + ) + + return result atleast_1d = partial(atleast_Nd, n=1) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 60643e2984..1186aeb35c 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -4364,7 +4364,8 @@ def test_atleast_Nd(): for n in range(1, 3): ary1, ary2 = dscalar(), dvector() - res_ary1, res_ary2 = atleast_Nd(ary1, ary2, n=n) + res_ary1 = atleast_Nd(ary1, n=n) + res_ary2 = atleast_Nd(ary2, n=n) assert res_ary1.ndim == n if n == ary2.ndim: From 6ead1f77bb174e574c43e50dd3168f5083741d96 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 12 Mar 2025 20:38:08 +0530 Subject: [PATCH 2/3] Modify atleast_Nd to accept only one positional argument --- tests/tensor/test_basic.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 1186aeb35c..8e80bf9399 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -4364,8 +4364,8 @@ def test_atleast_Nd(): for n in range(1, 3): ary1, ary2 = dscalar(), dvector() - res_ary1 = atleast_Nd(ary1, n=n) - res_ary2 = atleast_Nd(ary2, n=n) + res_ary1 = atleast_Nd(ary1, n) + res_ary2 = atleast_Nd(ary2, n) assert res_ary1.ndim == n if n == ary2.ndim: From b94687f12c164f063583f57ed501d24d1c47f656 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 12 Mar 2025 22:03:16 +0530 Subject: [PATCH 3/3] Modify atleast_Nd to have one array as the only positional argument --- pytensor/tensor/basic.py | 2 +- tests/tensor/test_basic.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index a6055824db..36eac9a461 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -4355,7 +4355,7 @@ def empty_like( def atleast_Nd( - arry: np.ndarray | TensorVariable, n: int = 1, left: bool = True + arry: np.ndarray | TensorVariable, *, n: int = 1, left: bool = True ) -> TensorVariable: """Convert input to an array with at least `n` dimensions.""" diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 8e80bf9399..1186aeb35c 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -4364,8 +4364,8 @@ def test_atleast_Nd(): for n in range(1, 3): ary1, ary2 = dscalar(), dvector() - res_ary1 = atleast_Nd(ary1, n) - res_ary2 = atleast_Nd(ary2, n) + res_ary1 = atleast_Nd(ary1, n=n) + res_ary2 = atleast_Nd(ary2, n=n) assert res_ary1.ndim == n if n == ary2.ndim: