I am observing different results using the tvm.relay.transform.FoldScaleAxis pass depending on the method I use to call the pass.
Calling the pass in the manner:
new_mod = relay.transform.FoldScaleAxis()(mod)
Differs from invoking it in the following manners (both of which produce equivalent results):
with tvm.transform.PassContext(opt_level=3):
new_mod = relay.transform.FoldScaleAxis()(mod)
seq = tvm.ir.transform.Sequential([relay.transform.FoldScaleAxis()])
with tvm.transform.PassContext(opt_level=3):
new_mod = seq(mod)
Is this expected behavior?
The following is example code demonstrating the behavior:
import tvm
from tvm import relay
import numpy as np
def create_model():
x = relay.var('x', shape=(1, 32, 10, 10), dtype='float32')
pool = relay.nn.adaptive_avg_pool2d(x, output_size=(1, 1))
relu = relay.nn.relu(pool)
mul = relay.multiply(x, relu)
conv = relay.nn.conv2d(mul, relay.const(np.random.rand(16, 32, 1, 1).astype('float32')), kernel_size=(1, 1), channels=16)
mod = tvm.IRModule.from_expr(conv)
return mod
mod = create_model()
mod_0 = relay.transform.FoldScaleAxis()(mod)
with tvm.transform.PassContext(opt_level=3):
mod_1 = relay.transform.FoldScaleAxis()(mod)
seq = tvm.ir.transform.Sequential([relay.transform.FoldScaleAxis()])
with tvm.transform.PassContext(opt_level=3):
mod_2 = seq(mod)
assert not tvm.ir.structural_equal(mod_0, mod_1)
assert tvm.ir.structural_equal(mod_1, mod_2)