Generating error checking and basic flow control

Linux-cpp-lisp: <p>Say I want to use FX to generate a <code>GraphModule</code> that should, as its first operation, run an assertion on the shape of its input argument. It is not possible <code>torch.fx.wrap</code> <code>assert</code>, since <code>assert</code> is not a function. Implementing it as an <code>if</code> is also not supported, since it counts as control flow.</p>
<p>I can implement an <code>asserter</code> function and <code>torch.fx.wrap</code> it, but that seems quite ugly and gives bad tracebacks:</p>
<pre><code class="lang-python">In [18]: def asserter(b: bool):
    ...:     assert b
    ...: 

In [19]: def f(x):
    ...:     asserter(x.shape[0] == 5)
    ...:     return x.sum()
    ...: 

In [20]: fx.wrap(asserter)
Out[20]: &lt;function __main__.asserter(b: bool)&gt;

In [21]: g = fx.symbolic_trace(f)

In [22]: print(g.code)
import __main__
def forward(self, x):
    getattr_1 = x.shape
    getitem = getattr_1[0];  getattr_1 = None
    eq_1 = getitem == 5;  getitem = None
    asserter = __main__.asserter(eq_1);  eq_1 = None
    sum_1 = x.sum();  x = None
    return sum_1
    

In [23]: s = torch.jit.script(g)

In [24]: s(torch.ones(7))

Error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "&lt;ipython-input-18-375f83c1af1b&gt;", line 2, in forward
def asserter(b: bool):
    assert b
    ~~~~~~~~ &lt;--- HERE
RuntimeError: AssertionError: 


In [25]: s(torch.ones(5))
Out[25]: tensor(5.)

In [26]: s.graph
Out[26]: 
graph(%self : __torch__.torch.fx.graph_module.___torch_mangle_0.f,
      %x.1 : Tensor):
  %17 : None = prim::Constant()
  %13 : Function = prim::Constant[name="asserter"]()
  %5 : int = prim::Constant[value=0]()
  %9 : int = prim::Constant[value=5]()
  %getattr_1.1 : int[] = aten::size(%x.1) # &lt;string&gt;:7:9
  %getitem.1 : int = aten::__getitem__(%getattr_1.1, %5)
  %eq_1.1 : bool = aten::eq(%getitem.1, %9)
  %asserter : None = prim::CallFunction(%13, %eq_1.1)
  %sum_1.1 : Tensor = aten::sum(%x.1, %17)
  return (%sum_1.1)
</code></pre>
<p>As you can see in the final output, it also complicates the compiled TorchScript significantly — using assert directly generates code without a Python function call:</p>
<pre><code class="lang-python">In [29]: def f2(x):
    ...:     assert x.shape[0] == 5
    ...:     return x.sum()
    ...: 

In [30]: s2 = torch.jit.script(f2)

In [31]: print(s2.graph)
graph(%x.1 : Tensor):
  %13 : None = prim::Constant()
  %24 : str = prim::Constant[value="AssertionError: "]()
  %3 : int = prim::Constant[value=0]() # &lt;ipython-input-29-fc36d8843bb5&gt;:2:19
  %5 : int = prim::Constant[value=5]() # &lt;ipython-input-29-fc36d8843bb5&gt;:2:25
  %2 : int[] = aten::size(%x.1) # &lt;string&gt;:7:9
  %4 : int = aten::__getitem__(%2, %3) # &lt;ipython-input-29-fc36d8843bb5&gt;:2:11
  %6 : bool = aten::eq(%4, %5) # &lt;ipython-input-29-fc36d8843bb5&gt;:2:11
   = prim::If(%6) # &lt;ipython-input-29-fc36d8843bb5&gt;:2:4
    block0():
      -&gt; ()
    block1():
       = prim::RaiseException(%24) # &lt;ipython-input-29-fc36d8843bb5&gt;:2:4
      -&gt; ()
  %14 : Tensor = aten::sum(%x.1, %13) # &lt;ipython-input-29-fc36d8843bb5&gt;:3:11
  return (%14)
</code></pre>
<p>Is there a better way to go about this?</p>
<p>Thanks!</p>
<p>(This is in some sense a subset of the discussion on <a href="https://discuss.pytorch.org/t/customizing-generated-code/114236/3" class="inline-onebox">Customizing generated code? - #3 by Linux-cpp-lisp</a>)</p>

jfix: <p><code>torch._assert()</code> is traceable, you should be able to use instead of builtin <code>assert</code>. Additionally, for modules where you may not want/be able to modify the source to use <code>torch._assert()</code>, there is an AST Rewriter tracer in FX experimental, see <a href="https://github.com/pytorch/pytorch/blob/06cb9293c572e86d5032eb1b36e18f7c00df1d74/torch/fx/experimental/rewriter.py#L68" rel="noopener nofollow ugc">here</a>.</p>

Linux-cpp-lisp: <p>Hi <a class="mention" href="/u/jfix">@jfix</a>,</p>
<p>Interesting — the AST rewriter looks like a great addition!</p>
<p><code>torch._assert</code> looks like a good interim solution. The only issue remaining with it is that it still makes the backtrace useless:</p>
<pre><code class="lang-python">In [4]: def f(x):
   ...:     torch._assert(x.shape[0] == 5, "invalid shape")
   ...:     return x*2.
   ...: 
   ...: 

In [5]: s = torch.jit.script(f)

In [6]: s(torch.ones(4))
---------------------------------------------------------------------------
Error                                     Traceback (most recent call last)
&lt;ipython-input-6-62f4f187df9a&gt; in &lt;module&gt;
----&gt; 1 s(torch.ones(4))

Error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "&lt;string&gt;", line 8, in f
  return out_size
def _assert(condition : bool, message : str):
  assert condition, message
  ~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
RuntimeError: AssertionError: invalid shape


In [7]: s(torch.ones(5))
Out[7]: tensor([2., 2., 2., 2., 2.])
</code></pre>
<p>This is a minor issue, though, and can be mitigated by using an assertion message that says where <code>torch._assert</code> was called.</p>