Using torch.fx.symbolic_trace with view(*shape)

lnair: <p>Hello,</p>
<p>I’m trying to use <code>torch.fx.symbolic_trace</code> to trace the <code>bert.SelfAttention</code> module. However, I’m running into the following error when <code>torch.Tensor.view</code> is called as <code>x = x.view(*new_shape)</code>:</p>
<p><code>torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when used in a for loop or as a *args or **kwargs function argument.</code></p>
<p>Possibly caused due to <code>*new_shape</code> argument to <code>view</code>? I’ve been trying to write a custom <code>Tracer</code> class with a new <code>iter</code> method, but I’m not able to resolve the issue. Is there a workaround to fixing this without modifying the actual self-attention code?</p>
<p>Full code:</p>
<pre><code class="lang-auto">from torch.fx import symbolic_trace
from transformers.models.bert import modeling_bert as bert
from transformers.models.bert import BertConfig

bert_config = BertConfig(hidden_size=10, num_attention_heads=5)
attention_module = bert.BertSelfAttention(bert_config)
graph = symbolic_trace(attention_module).graph
</code></pre>

tom: <p>Sometimes, these models can be involved, especially for a library with the genesis and target audience as transformers.</p>
<p>Chance has it that I released an initial commit of the <a href="https://github.com/mathinf/toroidal/" rel="noopener nofollow ugc">toroidal transformers</a> that has an attention module that works well with symbolic tracing:</p>
<pre><code class="lang-auto">m = toroidal.models.Attention(786, 3)
gr = torch.fx.symbolic_trace(m)
</code></pre>
<p>or in a BERT model (the multilingual from the test directory):</p>
<pre><code class="lang-auto">gr = torch.fx.symbolic_trace(model_bert.blocks[0].attn)

print(gr)
</code></pre>
<p>gives</p>
<pre><code class="lang-auto">Attention(
  (qkv): Linear(in_features=768, out_features=2304, bias=True)
  (dropout_attn): Dropout(p=0.1, inplace=False)
  (proj): Linear(in_features=768, out_features=768, bias=True)
  (dropout_out): Dropout(p=0.1, inplace=False)
)



def forward(self, x):
    getattr_1 = x.shape
    getitem = getattr_1[0]
    getitem_1 = getattr_1[1]
    getitem_2 = getattr_1[2];  getattr_1 = None
    qkv = self.qkv(x);  x = None
    view = qkv.view(getitem, getitem_1, 3, 12, -1);  qkv = None
    unbind = view.unbind(dim = 2);  view = None
    getitem_3 = unbind[0]
    getitem_4 = unbind[1]
    getitem_5 = unbind[2];  unbind = None
    einsum = torch.functional.einsum('bthc,bshc-&gt;bhts', getitem_3, getitem_4);  getitem_3 = getitem_4 = None
    mul = einsum * 0.125;  einsum = None
    softmax = mul.softmax(-1);  mul = None
    dropout_attn = self.dropout_attn(softmax);  softmax = None
    einsum_1 = torch.functional.einsum('bhts,bshc-&gt;bthc', dropout_attn, getitem_5);  dropout_attn = getitem_5 = None
    reshape = einsum_1.reshape(getitem, getitem_1, getitem_2);  einsum_1 = getitem = getitem_1 = getitem_2 = None
    proj = self.proj(reshape);  reshape = None
    dropout_out = self.dropout_out(proj);  proj = None
    return dropout_out
</code></pre>
<p>and I guess this is the disadvantage for manipulating the matmuls: I used einsum instead. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>If GPT works for you, you could check out <a href="https://github.com/karpathy/minGPT/" rel="noopener nofollow ugc">minGPT</a>, I’d imagine that A. Karpathy’s code might work more easily with the tracer.</p>
<p>Best regards</p>
<p>Thomas</p>

lnair: <p>Thanks for the quick response! I’ll look into this!</p>