Skip to content

Commit 4c52737

Browse files
committed
Fix indexing of at.stack in get_support_shape
1 parent 3a8d898 commit 4c52737

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

pymc/distributions/shape_utils.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -717,7 +717,7 @@ def get_support_shape(
717717
shape = to_tuple(shape)
718718
assert isinstance(shape, tuple)
719719
inferred_support_shape = at.stack(
720-
[shape[-i - 1] - support_shape_offset[-i - 1] for i in range(ndim_supp)]
720+
[shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)]
721721
)
722722

723723
if inferred_support_shape is None and dims is not None:
@@ -726,15 +726,15 @@ def get_support_shape(
726726
model = modelcontext(None)
727727
inferred_support_shape = at.stack(
728728
[
729-
model.dim_lengths[dims[-i - 1]] - support_shape_offset[-i - 1] # type: ignore
730-
for i in range(ndim_supp)
729+
model.dim_lengths[dims[i]] - support_shape_offset[i] # type: ignore
730+
for i in np.arange(-ndim_supp, 0)
731731
]
732732
)
733733

734734
if inferred_support_shape is None and observed is not None:
735735
observed = convert_observed_data(observed)
736736
inferred_support_shape = at.stack(
737-
[observed.shape[-i - 1] - support_shape_offset[-i - 1] for i in range(ndim_supp)]
737+
[observed.shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)]
738738
)
739739

740740
if inferred_support_shape is None:

0 commit comments

Comments
 (0)