Hi all,
Thanks for your work on this exciting new feature of PyTorch!
I’m interested in FX for an application that involves graph rewriting based on tensor shapes. The catch is that all of the shapes, except for a batch dimension, are known at “compile” time.
I can see one way of doing this with FX, using Transformer with real Tensors full of zeros and branching in call_function as demonstrated in the docs 8.
This seems a little inefficient, however: is it possible to attach (partial) shape information to a Proxy and propagate it through the graph, using it to rewrite some function calls in Transformer?
Thanks!