Workaround for unsupported return of multiple tensors for torch.cond in a model intended for torch.onnx.export(dynamo=true,...)

ionymikler: <p>Trying to <code>onnx.export</code> a <code>nn.Module</code> with a conditional in its computational graph. In essence similar to this example:</p>
<pre data-code-wrap="py"><code class="lang-py">import torch

class Wrapper(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cond_model = CondModel()
    
    def forward(self, x):
        nt = self.cond_model(x)
        return nt

class CondModel(torch.nn.Module):
    def forward(self, x):
        
        def true_fn(x,z):
            x = x + 1.0
            z = z * 0.0
            return x,z

        def false_fn(x,z):
            x = x - 1.0
            z = z * 1.0
            return x,z

        z = torch.rand(x.shape)
        nt = torch.cond(x.sum() &gt; 0, true_fn, false_fn, [x,z])
        return nt
</code></pre>
<p>As per <a href="https://pytorch.org/docs/2.6/cond.html#torch._higher_order_ops.cond.cond" rel="noopener nofollow ugc">the documentation</a>, the return from <code>torch.cond</code> must be a single tensor. Is there a dirty workaround that allows to get multiple tensors from the return?<br>
I tried using nested tensors:</p>
<pre data-code-wrap="py"><code class="lang-py">def true_fn(x,z):
   x = x + 1.0
   z = z * 0.0
   nt = torch.nested.nested_tensor([x,z], layout=torch.jagged)
   return nt
</code></pre>
<p>But compile fails at validation of the <code>.shape</code> of the return tensors (<code>.shape</code> in <code>NestedTensors</code> <a href="https://pytorch.org/docs/2.6/nested.html#size" rel="noopener nofollow ugc">loses precise meaning</a>):</p>
<pre data-code-wrap="pytb"><code class="lang-pytb">torch._dynamo.exc.Unsupported: Expect branches to return tensors with same metadata but find pair[0] differ in 'shape: torch.Size([2, s1]) vs torch.Size([2, s2])', 'stride: (s1, 1) vs (s2, 1)', where lhs is TensorMetadata(shape=torch.Size([2, s1]), dtype=torch.float32, requires_grad=False, stride=(s1, 1), memory_format=None, is_quantized=False, qparams={}) and rhs is TensorMetadata(shape=torch.Size([2, s2]), dtype=torch.float32, requires_grad=False, stride=(s2, 1), memory_format=None, is_quantized=False, qparams={})
</code></pre>
<details>
  **<summary>Full traceback here</summary>**
<pre data-code-wrap="pytb"><code class="lang-pytb">Traceback (most recent call last):
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 55, in graph_break_as_hard_error
  return fn(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 906, in call_function
  unimplemented(
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/exc.py", line 356, in unimplemented
  raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: Expect branches to return tensors with same metadata but find pair[0] differ in 'shape: torch.Size([2, s1]) vs torch.Size([2, s2])', 'stride: (s1, 1) vs (s2, 1)', where lhs is TensorMetadata(shape=torch.Size([2, s1]), dtype=torch.float32, requires_grad=False, stride=(s1, 1), memory_format=None, is_quantized=False, qparams={}) and rhs is TensorMetadata(shape=torch.Size([2, s2]), dtype=torch.float32, requires_grad=False, stride=(s2, 1), memory_format=None, is_quantized=False, qparams={})

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

Traceback (most recent call last):
File "/home/iony/DTU/f24/thesis/code/early_exit_vit/simple/example_conditional.py", line 34, in &lt;module&gt;
  result = model(input_tensor)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
  return forward_call(*args, **kwargs)
File "/home/iony/DTU/f24/thesis/code/early_exit_vit/simple/example_conditional.py", line 9, in forward
  nt = self.cond_model(x)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
  return forward_call(*args, **kwargs)
File "/home/iony/DTU/f24/thesis/code/early_exit_vit/simple/example_conditional.py", line 28, in forward
  nt = torch.cond(x.sum() &gt; 0, true_fn, false_fn, [x,z])
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_higher_order_ops/cond.py", line 201, in cond
  return torch.compile(_cond_op_wrapper, backend=backend, fullgraph=True)(
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 576, in _fn
  return fn(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 1406, in __call__
  return self._torchdynamo_orig_callable(
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 566, in __call__
  return _compile(
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 1006, in _compile
  guarded_code = compile_inner(code, one_graph, hooks, transform)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 734, in compile_inner
  return _compile_inner(code, one_graph, hooks, transform)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_utils_internal.py", line 95, in wrapper_function
  return function(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 769, in _compile_inner
  out_code = transform_code_object(code, transform)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 1402, in transform_code_object
  transformations(instructions, code_options)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 237, in _fn
  return fn(*args, **kwargs)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 681, in transform
  tracer.run()
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 2906, in run
  super().run()
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1076, in run
  while self.step():
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 986, in step
  self.dispatch_table[inst.opcode](self, inst)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 683, in wrapper
  return inner_fn(self, inst)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1763, in CALL_FUNCTION_EX
  self.call_function(fn, argsvars.items, kwargsvars)
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 921, in call_function
  self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 58, in graph_break_as_hard_error
  raise UncapturedHigherOrderOpError(reason + msg) from e
torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile. Scroll up to find out what causes the graph break.

from user code:
 File "/home/iony/miniconda3/envs/eevit/lib/python3.9/site-packages/torch/_higher_order_ops/cond.py", line 193, in _cond_op_wrapper
  return cond_op(*args, **kwargs)
</code></pre>
</details>
<p>Is the feature not implemented, even in a nightly? Or is there another workaround that might work if I intend to operate only for inference?</p>
<h3><a name="p-464177-error-logs-1" class="anchor" href="#p-464177-error-logs-1"></a>Error logs</h3>
<p><em>No response</em></p>
<h3><a name="p-464177-versions-2" class="anchor" href="#p-464177-versions-2"></a>Versions</h3>
<p>Versions of relevant libraries:<br>
[pip3] numpy==1.26.4<br>
[pip3] onnx==1.17.0<br>
[pip3] onnxruntime==1.19.2<br>
[pip3] onnxscript==0.1.0.dev20241226<br>
[pip3] torch==2.6.0.dev20241226+cu124</p>

ionymikler: <p>Solved <a href="https://github.com/pytorch/pytorch/issues/145734#issuecomment-2621257546" rel="noopener nofollow ugc">here</a></p>