Skip to content

Blockwise inner graph not cloned during function compilation #2028

@ricardoV94

Description

@ricardoV94

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions