Looks like tf.shape is not compatible with relay.vm . Instead of determining the shape at the runtime vm.compiler just creates new named parameter which should be provided by user when calling vm.run().
Example of TF2 model containing just one operator tf.shape:
import tensorflow as tf
import tvm
from tvm import relay
import numpy as np
from tvm.relay.frontend import tensorflow2
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
class MyModel2D(tf.Module):
def __init__(self):
super(MyModel2D, self).__init__()
@tf.function(input_signature=[tf.TensorSpec([None], tf.float32)])
def __call__(self, input0):
y=tf.shape(input0)
return {'output0': y}
m2d = MyModel2D()
f = m2d.__call__.get_concrete_function(tf.TensorSpec([None], tf.float32))
x=np.array([2,3,0,1,1,1,2], dtype="float32")
y=f(x)
print("TF res:", y)
ff = convert_variables_to_constants_v2(f, lower_control_flow=False)
graph_def = ff.graph.as_graph_def(add_shapes=True)
# Lets build relay module for particular input shape
mod, params = tensorflow2.from_tensorflow(graph_def, shape={"input0": [7]})
print("Parsing Done")
print(mod["main"])
ctx=tvm.cpu()
target="llvm"
with tvm.transform.PassContext(opt_level=3):
vm_exec = relay.vm.compile(mod, target=target) # params=params is missing !!!
print("Compilation Done")
vm = tvm.runtime.vm.VirtualMachine(vm_exec, ctx)
# Now I have to add additional named parameter Shape to the input data. Otherwise vm.run validation fails
ff_data = {"input0":x, "Shape":np.array([7], "int32")}
res = vm.run(**ff_data)
print("vm.run:", res)
Relay main function
fn (%Shape: Tensor[(1), int32]) {
%Shape
}
As you can see it does not use %input0 input at all. Instead of that it adds new param %Shape to main function args.