Fix in-place mutations from NUMBA on Blockwise inner graph rewrites #2232
Fix in-place mutations from NUMBA on Blockwise inner graph rewrites #2232Dekermanjian wants to merge 1 commit into
Conversation
…s on a clone of the inner graph, updated tests that expected the inner graph to be mutated in-place and added a new test to ensure inner graph rewrites occur on a clone
|
@ricardoV94 has to sign off on this. We talked about this over the weekend and he had some master plan to fix the issue upstream. I don't remember what it was :) |
|
Scan is supposed to be mutated in place, it was meant to be cloned when the function was created, but Blockwise hid it. We have to have a consistency logic and not force duplicate clones because we aren't using the logic consistently. My grand plan is ops being immutable but that can wait, for now we can work with current design . My suggestion now: have two classes of Blockwise, one with the HasInnerGraph mixin and make sure if we wrap Scan/OFG/MinimizeOps we use the new class. One of the methods the mixin demands is clone. This will be called on function creation and fix the bug and also give you dprint for the inner graph of Blockwise. (Or single Blockwise always with the mixin but no Op in clone if the inner Op is itself not HasInnerGraph) Numba dispatch should remain in place |
| fn_scan_op = fn_scan_node.op | ||
| assert isinstance(fn_scan_op, Scan) | ||
| fn_cholesky_op = fn_scan_op.fgraph.outputs[0].owner.op | ||
| opt_fgraph = numba_optimize_inner_fgraph(fn_scan_op, fn_scan_node) |
There was a problem hiding this comment.
not good this isn't confirming the final graph had it, it's confirming it could have
|
Thank you for the review and tips @ricardoV94. I will give this another pass following your suggestions. |
|
@Dekermanjian I'm trying the maximal solution let's see. Would suggest you wait a bit and see how it pans, maybe it aegues in favor of conservative cloning |
|
Okay, thank you for the heads up @ricardoV94. I will pause until I hear back from you. |
3 participants
Description
split
numba_funcify_Scaninto two functions in order to apply rewrites on a clone of the inner graph, updated tests that expected the inner graph to be mutated in-place and added a new test to ensure inner graph rewrites occur on a clone graph.Related Issue
Checklist
Type of change