Update pymc/logprob/utils.py

Co-authored-by: Ricardo Vieira <[email protected]>
pymc-devs · Dec 1, 2023 · 0c80a68 · 0c80a68
1 parent b8fde47
commit 0c80a68
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pymc/logprob/utils.py b/pymc/logprob/utils.py
@@ -212,7 +212,7 @@ def transform(self, node: Apply) -> Union[Variable, str, None]:
         """
         raise NotImplementedError
 
-    def rewrite(self, node: Union[Scan, OpFromGraph]):
+    def apply(self, fgraph: FunctionGraph, node: Union[Scan, OpFromGraph]):
         if not isinstance(node.op, (Scan, OpFromGraph)):
             raise TypeError("Expected Scan or OpFromGraph in InernerGraphRewriter")
         node_inputs, node_outputs = node.op.inner_inputs, node.op.inner_outputs