How to print shape inference in relay IR?

Questions
#1

Hi, i send a dynamic tensor to add op, and i want to see shape inference in relay IR, my code is:

x = relay.var("x", shape = (relay.Any(), 2), dtype = "float32")
y = relay.var("y", shape = (1, 2), dtype = "float32")

mod = tvm.IRModule()
mod["main"] = relay.Function([x, y], relay.add(x, y))

comp = relay.vm.VMCompiler()
opt_mod, _ = comp.optimize(mod, target = "llvm")
print(opt_mod.astext(show_meta_data = False))

and i get:

...
%0 = fn (%p0: Tensor[(?, 2), float32], %p1: Tensor[(1, 2), float32], Primitive=1) -> Tensor[(?, 2), float32] {
  add(%p0, %p1) /* ty=Tensor[(?, 2), float32] */
};
%1 = (%in_shape_0, meta[relay.Constant][1] /* ty=Tensor[(2), int64] */);
%2 = (%shape_func_out_0,);
let %shape_func: () = vm.shape_func(%0, %1, %2, meta[relay.attrs.ShapeFuncAttrs][0]) /* ty=() */;
...

How to get the detail of the %shape_func in relay IR? My TVM version is 0.8, thankyou very much!

#2

Hi Xiang,

As far as I know shape inference is a runtime behavior in Relay VM. And Add is a broadcast operator. So, I guess the shape_func of add is as same as broadcast that will be invoked at runtime.

#3

Hi, Thankyou for your reponse, i read TVM docs and know that shape functions will be fused in Relay. I want to see the contents of shape functions before and after fusion, is there any way? Thankyou very much!