Skip to content

Modify atleast_Nd to accept only one positional argument #1291

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 12, 2025
Merged

Modify atleast_Nd to accept only one positional argument #1291

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ce69b4d
Modify atleast_Nd to accept only one positional argument
Abhinav-Khot Mar 12, 2025
6ead1f7
Modify atleast_Nd to accept only one positional argument
Abhinav-Khot Mar 12, 2025
b94687f
Modify atleast_Nd to have one array as the only positional argument
Abhinav-Khot Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 12 additions & 18 deletions pytensor/tensor/basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to avoid surprises, I would make it so you have to specify n explicitly like before. It simply doesn't accept multiple arrays anymore, so it won't be a silent change if someone was relying on this.

Later on we can allow n to be positional only like you did now.

Suggested change
arry: np.ndarray | TensorVariable, n: int = 1, left: bool = True
arry: np.ndarray | TensorVariable, *, n: int = 1, left: bool = True

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have made this change, modified the tests accordingly too

) -> 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)
Expand Down
3 changes: 2 additions & 1 deletion tests/tensor/test_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
res_ary2 = atleast_Nd(ary2, n)

assert res_ary1.ndim == n
if n == ary2.ndim:
Expand Down