Add error message in Numba implementation of SpecifyShape

ricardoV94 · ricardoV94 · commit 0b2dcbe4fb83 · 2025-10-09T12:12:01.000+02:00
diff --git a/pytensor/link/numba/dispatch/basic.py b/pytensor/link/numba/dispatch/basic.py
@@ -545,11 +545,11 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):
     shape_input_names = ["shape_" + str(i) for i in range(len(shape_inputs))]
 
     func_conditions = [
-        f"assert x.shape[{i}] == {shape_input_names}"
-        for i, (shape_input, shape_input_names) in enumerate(
+        f"assert x.shape[{i}] == {eval_dim_name}, f'SpecifyShape: dim {{{i}}} of input has shape {{x.shape[{i}]}}, expected {{{eval_dim_name}.item()}}.'"
+        for i, (node_dim_input, eval_dim_name) in enumerate(
             zip(shape_inputs, shape_input_names, strict=True)
         )
-        if shape_input is not NoneConst
+        if node_dim_input is not NoneConst
     ]
 
     func = dedent(

Original file line number	Diff line number	Diff line change
`@@ -545,11 +545,11 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):`
`545`	`545`	`shape_input_names = ["shape_" + str(i) for i in range(len(shape_inputs))]`
`546`	`546`
`547`	`547`	`func_conditions = [`
`548`		`- f"assert x.shape[{i}] == {shape_input_names}"`
`549`		`- for i, (shape_input, shape_input_names) in enumerate(`
	`548`	`+ f"assert x.shape[{i}] == {eval_dim_name}, f'SpecifyShape: dim {{{i}}} of input has shape {{x.shape[{i}]}}, expected {{{eval_dim_name}.item()}}.'"`
	`549`	`+ for i, (node_dim_input, eval_dim_name) in enumerate(`
`550`	`550`	`zip(shape_inputs, shape_input_names, strict=True)`
`551`	`551`	`)`
`552`		`- if shape_input is not NoneConst`
	`552`	`+ if node_dim_input is not NoneConst`
`553`	`553`	`]`
`554`	`554`
`555`	`555`	`func = dedent(`