Can we add shape check when rewrite the computation graph in torchScript?

I want to rewrite the computation graph, transform the origin op to custom op, but my custom op has tensor shape restriction, can we add shape check when rewrite?

Can you share a minimal reproducible example so people can debug your problem?

1 Like
pattern = """ 
replacement = """ 
torch._C._jit_pass_custom_pattern_based_rewrite_graph(pattern, replacement,graph)

For operation index_mul::index_mul(%x,%2,%1), it wants %x to be 2 dim tensor, %2: 2 dim tensor, %1: 1 dim tensor, can we add shape check in subgraph rewrite?