Skip to content

Commit c62d733

Browse files
committed
Add error message in Numba implementation of SpecifyShape
1 parent 81988be commit c62d733

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

pytensor/link/numba/dispatch/basic.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -481,11 +481,11 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):
481481
shape_input_names = ["shape_" + str(i) for i in range(len(shape_inputs))]
482482

483483
func_conditions = [
484-
f"assert x.shape[{i}] == {shape_input_names}"
485-
for i, (shape_input, shape_input_names) in enumerate(
484+
f"assert x.shape[{i}] == {eval_dim_name}, f'SpecifyShape: dim {{{i}}} of input has shape {{x.shape[{i}]}}, expected {{{eval_dim_name}.item()}}.'"
485+
for i, (node_dim_input, eval_dim_name) in enumerate(
486486
zip(shape_inputs, shape_input_names, strict=True)
487487
)
488-
if shape_input is not NoneConst
488+
if node_dim_input is not NoneConst
489489
]
490490

491491
func = dedent(

0 commit comments

Comments
 (0)