Description
Preface: Inner graphs should never be mutable, we will get closer to it after #1908
This showed up in pymc-devs/pymc#7380
We were creating a graph that had a Blockwise whose inner op is an OFG (an AllocDiag). At the start of function compilation we clone the whole graph and deepclone ops with inner graphs, but this is done by asking if the op has the HasInnerGraphOp parentclass/mixin. Blockwise (which may or may not have one), doesn't have such a baseclass (It's clone method isn't handling internal clone either).
So long-story short, the test was compiling it once with numba, which introduced an inplace operation. And then another time with python which complained about it.
Why doesn't numba clone the OpFromGraph before mutating it? Because it was supposed to already have been cloned, and it would be wasteful to deep clone everything twice.
import pytensor
import pytensor.tensor as pt
import numpy as np
from pytensor.compile.builders import OpFromGraph
from pytensor.tensor.blockwise import Blockwise
x = pt.vector("x", shape=(3,))
# Blockwise wrapping an OpFromGraph (AllocDiag) inside an outer OpFromGraph
# d = pt.vectorize(pt.diag, signature="(n)->(n,n)")(x)
d = Blockwise(pt.diag(x).owner.op, signature="(n)->(n,n)")(x)
ofg = OpFromGraph([x], [d])
# Compiling with NUMBA mutates AllocDiag's inner fgraph because
# Blockwise is not HasInnerGraph, so graph deep copy doesn't recurse into its core_op.
a = pt.vector("a", shape=(3,))
pytensor.function([a], ofg(a), mode="NUMBA")
# AllocDiag's lazy perform compilation now fails because its
# fgraph was mutated with inplace ops by the numba rewrites.
b = pt.vector("b", shape=(3,))
f = pytensor.function([b], ofg(b), mode="FAST_COMPILE")
f(np.array([4.0, 5.0, 6.0]))
# TypeError: Graph must not contain inplace operations
I think the best course of action is to wait for #1908 and then add this as a regression test. OpFromGraph inner graph will not be mutable then.
Description
Preface: Inner graphs should never be mutable, we will get closer to it after #1908
This showed up in pymc-devs/pymc#7380
We were creating a graph that had a Blockwise whose inner op is an OFG (an AllocDiag). At the start of function compilation we clone the whole graph and deepclone ops with inner graphs, but this is done by asking if the op has the
HasInnerGraphOpparentclass/mixin. Blockwise (which may or may not have one), doesn't have such a baseclass (It's clone method isn't handling internal clone either).So long-story short, the test was compiling it once with numba, which introduced an inplace operation. And then another time with python which complained about it.
Why doesn't numba clone the OpFromGraph before mutating it? Because it was supposed to already have been cloned, and it would be wasteful to deep clone everything twice.
I think the best course of action is to wait for #1908 and then add this as a regression test. OpFromGraph inner graph will not be mutable then.