Error: torch.jit.trace with example_kwarg_inputs and nested dictionaries

username: shirgur
cooked: <p>Hi,</p>
<p>First, I know there is an issue with nested dictionaries when using “example_inputs”, but I was hoping it is not the same with “example_kwarg_imputs”.</p>
<p>Below are three toy examples for reproducing the error (and a “good” example").<br>
Bottom line - how can I trace a nested dictionary with “example_kwarg_imputs”?</p>
<h1><a name="example-1-works-1" class="anchor" href="#example-1-works-1"></a>Example 1 (works)</h1>
<h2><a name="code-2" class="anchor" href="#code-2"></a>Code</h2>
<pre><code class="lang-python">import torch

def func(a, b):
    return a + b

x = {"a": torch.tensor(1), "b": torch.tensor(2)}

ei_traced = torch.jit.trace(func, example_inputs=(x["a"], x["b"]))
print(ei_traced.graph)

eki_traced = torch.jit.trace(func, example_kwarg_inputs=x)
print(eki_traced.graph)
</code></pre>
<h2><a name="output-3" class="anchor" href="#output-3"></a>Output</h2>
<pre><code class="lang-auto">graph(%a : Long(requires_grad=0, device=cpu),
      %b : Long(requires_grad=0, device=cpu)):
  %2 : int = prim::Constant[value=1]() # &lt;workdir&gt;/test.py:4:0
  %3 : Long(requires_grad=0, device=cpu) = aten::add(%a, %b, %2) # &lt;workdir&gt;/test.py:4:0
  return (%3)

graph(%a : Long(requires_grad=0, device=cpu),
      %b : Long(requires_grad=0, device=cpu)):
  %2 : int = prim::Constant[value=1]() # &lt;workdir&gt;/test.py:4:0
  %3 : Long(requires_grad=0, device=cpu) = aten::add(%a, %b, %2) # &lt;workdir&gt;/test.py:4:0
  return (%3)
</code></pre>
<h1><a name="example-2-error-4" class="anchor" href="#example-2-error-4"></a>Example 2 (Error)</h1>
<h2><a name="code-5" class="anchor" href="#code-5"></a>Code</h2>
<pre><code class="lang-python">import torch

def func(a, b):
    return a["a2"] + b

x = {"a": {"a2": torch.tensor(1)}, "b": torch.tensor(2)}

ei_traced = torch.jit.trace(func, example_inputs=(x["a"], x["b"]))
print(ei_traced.graph)

eki_traced = torch.jit.trace(func, example_kwarg_inputs=x)
print(eki_traced.graph)
</code></pre>
<h2><a name="output-6" class="anchor" href="#output-6"></a>Output</h2>
<pre><code class="lang-python">graph(%a : Dict(str, Tensor),
      %b : Long(requires_grad=0, device=cpu)):
  %1 : str = prim::Constant[value="a2"]()
  %2 : Tensor = aten::__getitem__(%a, %1) # &lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py:859:0
  %4 : int = prim::Constant[value=1]() # &lt;workdir&gt;/test.py:4:0
  %5 : Long(requires_grad=0, device=cpu) = aten::add(%2, %b, %4) # &lt;workdir&gt;/test.py:4:0
  return (%5)

Traceback (most recent call last):
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 459, in run_mod_and_filter_tensor_outputs
    outs = wrap_retval(mod(**inputs))
RuntimeError: func() expected at most 1 argument(s) but received 2 argument(s). Declaration: func(Tensor a) -&gt; Tensor

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "&lt;workdir&gt;/test.py", line 11, in &lt;module&gt;
    eki_traced = torch.jit.trace(func, example_kwarg_inputs=x)
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 884, in trace
    _check_trace(
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 552, in _check_trace
    traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace")
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 467, in run_mod_and_filter_tensor_outputs
    raise TracingCheckError(
torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
encountered an exception while running the trace with test inputs.
Exception:
        func() expected at most 1 argument(s) but received 2 argument(s). Declaration: func(Tensor a) -&gt; Tensor
</code></pre>
<h1><a name="example-3-error-7" class="anchor" href="#example-3-error-7"></a>Example 3 (Error)</h1>
<h2><a name="code-8" class="anchor" href="#code-8"></a>Code</h2>
<pre><code class="lang-python">import torch

def func(a, b):
    return a["a2"] + b["b2"]

x = {"a": {"a2": torch.tensor(1)}, "b": {"b2": torch.tensor(2)}}

ei_traced = torch.jit.trace(func, example_inputs=(x["a"], x["b"]))
print(ei_traced.graph)

eki_traced = torch.jit.trace(func, example_kwarg_inputs=x)
print(eki_traced.graph)
</code></pre>
<h2><a name="output-9" class="anchor" href="#output-9"></a>Output</h2>
<pre><code class="lang-python">graph(%a : Dict(str, Tensor),
      %b : Dict(str, Tensor)):
  %1 : str = prim::Constant[value="a2"]()
  %2 : Tensor = aten::__getitem__(%a, %1) # &lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py:859:0
  %4 : str = prim::Constant[value="b2"]()
  %5 : Tensor = aten::__getitem__(%b, %4) # &lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py:859:0
  %6 : int = prim::Constant[value=1]() # &lt;workdir&gt;/test.py:4:0
  %7 : Long(requires_grad=0, device=cpu) = aten::add(%2, %5, %6) # &lt;workdir&gt;/test.py:4:0
  return (%7)

Traceback (most recent call last):
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 459, in run_mod_and_filter_tensor_outputs
    outs = wrap_retval(mod(**inputs))
RuntimeError: func() expected at most 0 argument(s) but received 2 argument(s). Declaration: func() -&gt; Tensor

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "&lt;workdir&gt;/test.py", line 11, in &lt;module&gt;
    eki_traced = torch.jit.trace(func, example_kwarg_inputs=x)
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 884, in trace
    _check_trace(
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 552, in _check_trace
    traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace")
  File "&lt;conda_env&gt;/lib/python3.8/site-packages/torch/jit/_trace.py", line 467, in run_mod_and_filter_tensor_outputs
    raise TracingCheckError(
torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
encountered an exception while running the trace with test inputs.
Exception:
        func() expected at most 0 argument(s) but received 2 argument(s). Declaration: func() -&gt; Tensor
</code></pre>