Hi @jfix,
Interesting — the AST rewriter looks like a great addition!
torch._assert looks like a good interim solution. The only issue remaining with it is that it still makes the backtrace useless:
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)
<ipython-input-6-62f4f187df9a> in <module>
----> 1 s(torch.ones(4))

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


In [7]: s(torch.ones(5))
Out[7]: tensor([2., 2., 2., 2., 2.])

This is a minor issue, though, and can be mitigated by using an assertion message that says where torch._assert was called.