Add missing argument to ShapeFeature.get_shape

brandonwillard · brandonwillard · commit 4134b22007f1 · 2023-06-28T19:38:38.000-04:00
diff --git a/aeppl/printing.py b/aeppl/printing.py
@@ -380,7 +380,7 @@ def process_shape_info(cls, output: Variable, pstate: Optional[PrinterStateType]
                 try:
                     old_precedence = getattr(pstate, "precedence", None)
                     pstate.precedence = new_precedence
-                    _s_i_out = shape_feature.get_shape(output, i)
+                    _s_i_out = shape_feature.get_shape(pstate.fgraph, output, i)
 
                     if not isinstance(_s_i_out, (Constant, TensorVariable)):
                         s_i_out = pstate.pprinter.process(_s_i_out, pstate)
diff --git a/aeppl/scan.py b/aeppl/scan.py
@@ -439,7 +439,8 @@ def update_scan_value_vars(
         # graph, so we use the shape feature to (hopefully) get the shape
         # without the entire `Scan` itself.
         full_out_shape = tuple(
-            fgraph.shape_feature.get_shape(full_out, i) for i in range(full_out.ndim)
+            fgraph.shape_feature.get_shape(fgraph, full_out, i)
+            for i in range(full_out.ndim)
         )
         new_val_var = at.empty(full_out_shape, dtype=full_out.dtype)
 

Original file line number	Diff line number	Diff line change
`@@ -439,7 +439,8 @@ def update_scan_value_vars(`
`439`	`439`	`# graph, so we use the shape feature to (hopefully) get the shape`
`440`	`440`	# without the entire `Scan` itself.
`441`	`441`	`full_out_shape = tuple(`
`442`		`- fgraph.shape_feature.get_shape(full_out, i) for i in range(full_out.ndim)`
	`442`	`+ fgraph.shape_feature.get_shape(fgraph, full_out, i)`
	`443`	`+ for i in range(full_out.ndim)`
`443`	`444`	`)`
`444`	`445`	`new_val_var = at.empty(full_out_shape, dtype=full_out.dtype)`
`445`	`446`