for output, variable in zip(outputs, variables): output[0] = variable @local_optimizer([OpFromGraph]) def inline_ofg_expansion(fgraph, node): """ This optimization expands internal graph of OpFromGraph. Only performed if node.op.is_inline == True Doing so can improve optimization at the cost of compilation speed. """ op = node.op if not isinstance(op, OpFromGraph): return False if not op.is_inline: return False return clone_replace(op.inner_outputs, {u: v for u, v in zip(op.inner_inputs, node.inputs)}) # We want to run this before the first merge optimizer # and before the first scan optimizer. optdb.register( "inline_ofg_expansion", in2out(inline_ofg_expansion), "fast_compile", "fast_run", position=-0.01, )
def inline_ofg_expansion(node): """ This optimization expands internal graph of OpFromGraph. Only performed if node.op.is_inline == True Doing so can improve optimization at the cost of compilation speed. """ op = node.op if not isinstance(op, OpFromGraph): return False if not op.is_inline: return False return aesara.clone( op.local_outputs, {u: v for u, v in zip(node.op.local_inputs, node.inputs)}) # We want to run this before the first merge optimizer # and before the first scan optimizer. optdb.register( "inline_ofg_expansion", gof.opt.in2out(inline_ofg_expansion), -0.01, "fast_compile", "fast_run", ) # Since OpFromGraph contains a Aesara compiled function, # we should let DebugMode know about it ops_with_inner_function[OpFromGraph] = "fn"