FX `Proxy` and shapes

Linux-cpp-lisp: <p>Hi all,</p>
<p>Thanks for your work on this exciting new feature of PyTorch!</p>
<p>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.</p>
<p>I can see one way of doing this with FX, using <code>Transformer</code> with real <code>Tensor</code>s full of zeros and branching in <code>call_function</code> as demonstrated in <a href="https://pytorch.org/docs/master/fx.html#torch.fx.Transformer" rel="noopener nofollow ugc">the docs</a>.</p>
<p>This seems a little inefficient, however: is it possible to attach (partial) shape information to a <code>Proxy</code> and propagate it through the graph, using it to rewrite some function calls in <code>Transformer</code>?</p>
<p>Thanks!</p>

Linux-cpp-lisp: <p>In a similar vein, if I do use <code>ShapeProp</code> to follow dummy <code>Tensor</code>s through the graph, is there a way to keep track of which dimensions “are” or derive from the batch dimension? (Short of picking a unique, and thus likely large, number?)</p>

James_Reed: <p>Hi <a class="mention" href="/u/linux-cpp-lisp">@Linux-cpp-lisp</a>,</p>
<p>This is a great question and one that’s come up several times. The long-and-short of it is that providing this sort of symbolic shape functionality is quite difficult, since it involves writing code or symbolic formulae to calculate the shapes for every op/module you might see in the program. We currently don’t have such formulae for general usage in PyTorch, though we might sometime in the future.</p>
<p>If, however, your transformation can make a simplifying assumption, it may be more tractable to implement for your use case. For example, if your transform only operates on a small set of operators or modules, it may be tenable to implement the formulas for just those ops/modules.</p>
<p>Please let me know if that works for your situation, otherwise we can register this need with the team to better motivate the general shape analysis problem.</p>
<p>James</p>

Linux-cpp-lisp: <p>Hi <a class="mention" href="/u/james_reed">@James_Reed</a> ,</p>
<p>Thanks for the quick response!</p>
<p>That makes sense — things seem to have worked out for my specific instance of the problem I was facing such that <code>ShapeProp</code> worked well enough, but for the general case one would definitely need symbolic shape tracing, particularly for keeping track of <em>which</em> dimensions of intermediate results are affected by the dimensions of the inputs.</p>
<p>(So its not really an issue of the efficiency of shape propagation with real <code>Tensor</code>s — it’s a more fundamental limitation, in general.)</p>
<p>I’m sure that general symbolic shape analysis for a library the size of PyTorch would be a <em>very</em> painful project of dubious general value. What might be interesting in the future is providing a common repository of symbolic shape formulas somewhere in <code>fx</code> that people can contribute back to over time — basically a small but growing <code>dict</code> of rules an an associated example of a <code>SymbolicShapeProp</code> class. The biggest challenge, as far as I can tell, would be a symbolic version of broadcasting. Once you know how two shapes broadcast, most other shape transformations that I can think of follow simple rules. (Einsum or matmul, for example.)</p>
<p>Anyway, that’s not a feature request, just some thoughts. Really appreciate the work you’ve all been doing on <code>fx</code> — right before it was released and I learned about it I was writing my own very ad-hoc limited version of it, and having this has really simplified things.</p>
<p>Best,<br>
A</p>

Chillee: <p>One idea I’ve been toying around with is some kind of general solver (i.e. SMT)-based shape inference. It’s unfortunately not that easy to propagate shapes in the general case - broadcasting is one significant example. However, there are lots of other annoying examples as well.</p>
<p>For example, <code>torch.cat([a, b])</code> results in a shape that’s dynamic with the shapes of both <code>a</code> and <code>b</code>.</p>
<p>A solver-based shape inference would allow you to query whether a shape 1. could be dynamic, and 2. if it were dynamic, what the constraints on its shape would be.</p>

Eugenia: <p>I am so happy to see that the issue have been solved Looking for the same info.</p>