Skip to content
Closed
Changes from 2 commits
Commits
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
4 changes: 2 additions & 2 deletions pymc/distributions/shape_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -472,8 +472,8 @@ def convert_size(size: Size) -> Optional[StrongSize]:
if size is None:
return None

if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim == 0):
size = (size,)
if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim==1 or size.ndim==0):
size = tuple(size,)
elif isinstance(size, (list, tuple)):
size = tuple(size)
else:
Copy link
Member

Choose a reason for hiding this comment

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

Does this work? I don't think we have to do anything if it is already a TensorVariable. We could assert it has ndim < 2. but that's about it.

Suggested change
if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim==1 or size.ndim==0):
size = tuple(size,)
elif isinstance(size, (list, tuple)):
size = tuple(size)
else:
if isinstance(size, int):
size = tuple(size,)
elif isinstance(size, (list, tuple)):
size = tuple(size)
elif not isinstance(TensorVariable):

Expand Down