[ONNX] Update API to torch.onnx.export(..., dynamo=True)

[titaiwangms]

Fixes #3200

Description

Checklist

• The issue that is being fixed is referred in the description (see above "Fixes #ISSUE_NUMBER")
• Only one issue is addressed in this pull request
• Labels from the issue that this PR is fixing are added to this pull request
• No unnecessary issues are included into this pull request.

TODO: The registry doc also needs an update.

cc @justinchuby

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/3223

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[titaiwangms]

@svekars PTAL

[justinchuby]

We should include the result of the print here I think

[titaiwangms]

What do you mean by result of the print? More than model proto?

[justinchuby]

Like this where we include the print output: https://github.com/pytorch/tutorials/pull/3223/files#diff-974479b5e9f94c102ca2630bcc39ee4e63d28f220dc1cea6198fc0083396a2bcR107-R135

[xadupre]

I looked at other examples in the tutorials, output was not included. The generation should take care of that otherwise we would have to update them everytime pytorch is released. I guess they do something very similar to what sphinx-gallery does.

[titaiwangms]

Yes, it should be printed in the page. The code will be executed.

[justinchuby]

I see. Strange I was looking at the registry page and there was no print out

[justinchuby]

I was looking at https://pytorch.org/tutorials/beginner/onnx/onnx_registry_tutorial.html

[justinchuby]

<
    ir_version=10,
    opset_imports={'pkg.onnxscript.torch_lib.common': 1, '': 18},
    producer_name='pytorch',
    producer_version='2.6.0',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,
    inputs=(
        %"input_1"<FLOAT,[3]>
    ),
    outputs=(
        %"mul"<FLOAT,[1]>
    ),
    initializers=(
        %"model.mlp.0.bias"<FLOAT,[2]>,
        %"model.mlp.0.weight"<FLOAT,[2,3]>,
        %"model.mlp.1.bias"<FLOAT,[1]>,
        %"model.mlp.1.weight"<FLOAT,[1,2]>
    ),
) {
    0 |  # node_Transpose_0
         %"val_0"<?,?> ⬅️ ::Transpose(%"model.mlp.0.weight") {perm=[1, 0]}
    1 |  # node_MatMul_1
         %"val_1"<?,?> ⬅️ ::MatMul(%"input_1", %"val_0")
    2 |  # node_Add_2
         %"linear"<FLOAT,[2]> ⬅️ ::Add(%"val_1", %"model.mlp.0.bias")
    3 |  # node_Transpose_3
         %"val_2"<?,?> ⬅️ ::Transpose(%"model.mlp.1.weight") {perm=[1, 0]}
    4 |  # node_MatMul_4
         %"val_3"<?,?> ⬅️ ::MatMul(%"linear", %"val_2")
    5 |  # node_Add_5
         %"linear_1"<FLOAT,[1]> ⬅️ ::Add(%"val_3", %"model.mlp.1.bias")
    6 |  # node_Constant_6
         %"val_4"<?,?> ⬅️ ::Constant() {value=Tensor<INT64,[]>(array(2), name=None)}
    7 |  # node_Cast_7
         %"scalar_tensor_default"<FLOAT,[]> ⬅️ ::Cast(%"val_4") {to=FLOAT}
    8 |  # node_Mul_8
         %"mul"<FLOAT,[1]> ⬅️ ::Mul(%"linear_1", %"scalar_tensor_default")
    return %"mul"<FLOAT,[1]>
}

[justinchuby]

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_mlp_0_weight: "f32[2, 3]", p_mlp_0_bias: "f32[2]", p_mlp_1_weight: "f32[1, 2]", p_mlp_1_bias: "f32[1]", x: "f32[3]"):
            linear: "f32[2]" = torch.ops.aten.linear.default(x, p_mlp_0_weight, p_mlp_0_bias);  x = p_mlp_0_weight = p_mlp_0_bias = None
            linear_1: "f32[1]" = torch.ops.aten.linear.default(linear, p_mlp_1_weight, p_mlp_1_bias);  linear = p_mlp_1_weight = p_mlp_1_bias = None
            
            sum_1: "f32[]" = torch.ops.aten.sum.default(linear_1)
            gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None
            
            true_graph_0 = self.true_graph_0
            false_graph_0 = self.false_graph_0
            cond = torch.ops.higher_order.cond(gt, true_graph_0, false_graph_0, [linear_1]);  gt = true_graph_0 = false_graph_0 = linear_1 = None
            getitem: "f32[1]" = cond[0];  cond = None
            return (getitem,)
            
        class true_graph_0(torch.nn.Module):
            def forward(self, linear_1: "f32[1]"):
                mul: "f32[1]" = torch.ops.aten.mul.Tensor(linear_1, 2);  linear_1 = None
                return (mul,)
                
        class false_graph_0(torch.nn.Module):
            def forward(self, linear_1: "f32[1]"):
                neg: "f32[1]" = torch.ops.aten.neg.default(linear_1);  linear_1 = None
                return (neg,)

[justinchuby]

<
    ir_version=10,
    opset_imports={'pkg.onnxscript.torch_lib.common': 1, '': 18, 'pkg.torch.__subgraph__': 1},
    producer_name='pytorch',
    producer_version='2.6.0',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,
    inputs=(
        %"x"<FLOAT,[3]>
    ),
    outputs=(
        %"getitem"<FLOAT,[1]>
    ),
    initializers=(
        %"mlp.0.bias"<FLOAT,[2]>,
        %"mlp.1.bias"<FLOAT,[1]>
    ),
) {
     0 |  # node_Constant_11
          %"val_0"<FLOAT,[3,2]> ⬅️ ::Constant() {value=Tensor<FLOAT,[3,2]>(array([[ 0.32409453,  0.09968598],
                 [ 0.23967852, -0.04969374],
                 [-0.09462868,  0.34749857]], dtype=float32), name='val_0')}
     1 |  # node_MatMul_1
          %"val_1"<FLOAT,[2]> ⬅️ ::MatMul(%"x", %"val_0")
     2 |  # node_Add_2
          %"linear"<FLOAT,[2]> ⬅️ ::Add(%"val_1", %"mlp.0.bias")
     3 |  # node_Constant_12
          %"val_2"<FLOAT,[2,1]> ⬅️ ::Constant() {value=Tensor<FLOAT,[2,1]>(array([[0.19137527],
                 [0.29681835]], dtype=float32), name='val_2')}
     4 |  # node_MatMul_4
          %"val_3"<FLOAT,[1]> ⬅️ ::MatMul(%"linear", %"val_2")
     5 |  # node_Add_5
          %"linear_1"<FLOAT,[1]> ⬅️ ::Add(%"val_3", %"mlp.1.bias")
     6 |  # node_ReduceSum_6
          %"sum_1"<FLOAT,[]> ⬅️ ::ReduceSum(%"linear_1") {noop_with_empty_axes=0, keepdims=False}
     7 |  # node_Constant_13
          %"scalar_tensor_default"<FLOAT,[]> ⬅️ ::Constant() {value=Tensor<FLOAT,[]>(array(0., dtype=float32), name='scalar_tensor_default')}
     8 |  # node_Greater_9
          %"gt"<BOOL,[]> ⬅️ ::Greater(%"sum_1", %"scalar_tensor_default")
     9 |  # node_If_10
          %"getitem"<FLOAT,[1]> ⬅️ ::If(%"gt") {then_branch=
              graph(
                  name=true_graph_0,
                  inputs=(

                  ),
                  outputs=(
                      %"mul_true_graph_0"<FLOAT,[1]>
                  ),
              ) {
                  0 |  # node_Constant_1
                       %"scalar_tensor_default_2"<FLOAT,[]> ⬅️ ::Constant() {value=Tensor<FLOAT,[]>(array(2., dtype=float32), name='scalar_tensor_default_2')}
                  1 |  # node_Mul_2
                       %"mul_true_graph_0"<FLOAT,[1]> ⬅️ ::Mul(%"linear_1", %"scalar_tensor_default_2")
                  return %"mul_true_graph_0"<FLOAT,[1]>
              }, else_branch=
              graph(
                  name=false_graph_0,
                  inputs=(

                  ),
                  outputs=(
                      %"neg_false_graph_0"<FLOAT,[1]>
                  ),
              ) {
                  0 |  # node_Neg_0
                       %"neg_false_graph_0"<FLOAT,[1]> ⬅️ ::Neg(%"linear_1")
                  return %"neg_false_graph_0"<FLOAT,[1]>
              }}
    return %"getitem"<FLOAT,[1]>
}

[justinchuby]

<
    ir_version=10,
    opset_imports={'pkg.onnxscript.torch_lib.common': 1, '': 18, 'pkg.torch.__subgraph__': 1},
    producer_name='pytorch',
    producer_version='2.6.0',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,
    inputs=(
        %"x"<FLOAT,[3]>
    ),
    outputs=(
        %"getitem"<FLOAT,[1]>
    ),
    initializers=(
        %"mlp.0.weight"<FLOAT,[2,3]>,
        %"mlp.0.bias"<FLOAT,[2]>,
        %"mlp.1.weight"<FLOAT,[1,2]>,
        %"mlp.1.bias"<FLOAT,[1]>
    ),
) {
     0 |  # node_Transpose_0
          %"val_0"<?,?> ⬅️ ::Transpose(%"mlp.0.weight") {perm=[1, 0]}
     1 |  # node_MatMul_1
          %"val_1"<?,?> ⬅️ ::MatMul(%"x", %"val_0")
     2 |  # node_Add_2
          %"linear"<FLOAT,[2]> ⬅️ ::Add(%"val_1", %"mlp.0.bias")
     3 |  # node_Transpose_3
          %"val_2"<?,?> ⬅️ ::Transpose(%"mlp.1.weight") {perm=[1, 0]}
     4 |  # node_MatMul_4
          %"val_3"<?,?> ⬅️ ::MatMul(%"linear", %"val_2")
     5 |  # node_Add_5
          %"linear_1"<FLOAT,[1]> ⬅️ ::Add(%"val_3", %"mlp.1.bias")
     6 |  # node_ReduceSum_6
          %"sum_1"<FLOAT,[]> ⬅️ ::ReduceSum(%"linear_1") {noop_with_empty_axes=0, keepdims=False}
     7 |  # node_Constant_7
          %"val_4"<?,?> ⬅️ ::Constant() {value=Tensor<INT64,[]>(array(0), name=None)}
     8 |  # node_Cast_8
          %"scalar_tensor_default"<FLOAT,[]> ⬅️ ::Cast(%"val_4") {to=FLOAT}
     9 |  # node_Greater_9
          %"gt"<BOOL,[]> ⬅️ ::Greater(%"sum_1", %"scalar_tensor_default")
    10 |  # node_If_10
          %"getitem"<FLOAT,[1]> ⬅️ ::If(%"gt") {then_branch=
              graph(
                  name=true_graph_0,
                  inputs=(

                  ),
                  outputs=(
                      %"mul_true_graph_0"<?,?>
                  ),
              ) {
                  0 |  # node_true_graph_0_0
                       %"mul_true_graph_0"<?,?> ⬅️ pkg.torch.__subgraph__::true_graph_0(%"linear_1")
                  return %"mul_true_graph_0"<?,?>
              }, else_branch=
              graph(
                  name=false_graph_0,
                  inputs=(

                  ),
                  outputs=(
                      %"neg_false_graph_0"<?,?>
                  ),
              ) {
                  0 |  # node_false_graph_0_0
                       %"neg_false_graph_0"<?,?> ⬅️ pkg.torch.__subgraph__::false_graph_0(%"linear_1")
                  return %"neg_false_graph_0"<?,?>
              }}
    return %"getitem"<FLOAT,[1]>
}

<
    opset_imports={'': 18},
>
def pkg.torch.__subgraph__::false_graph_0(
    inputs=(
        %"linear_1"<FLOAT,[1]>
    ),
    outputs=(
        %"neg"<FLOAT,[1]>
    ),
) {
    0 |  # node_Neg_0
         %"neg"<FLOAT,[1]> ⬅️ ::Neg(%"linear_1")
    return %"neg"<FLOAT,[1]>
}

<
    opset_imports={'': 18},
>
def pkg.torch.__subgraph__::true_graph_0(
    inputs=(
        %"linear_1"<FLOAT,[1]>
    ),
    outputs=(
        %"mul"<FLOAT,[1]>
    ),
) {
    0 |  # node_Constant_0
         %"val_0"<?,?> ⬅️ ::Constant() {value=Tensor<INT64,[]>(array(2), name=None)}
    1 |  # node_Cast_1
         %"scalar_tensor_default"<FLOAT,[]> ⬅️ ::Cast(%"val_0") {to=FLOAT}
    2 |  # node_Mul_2
         %"mul"<FLOAT,[1]> ⬅️ ::Mul(%"linear_1", %"scalar_tensor_default")
    return %"mul"<FLOAT,[1]>
}

[justinchuby]

Suggested change

	# producer_version='2.7.0.dev20250124+cu124',
	# producer_version='2.6.0',

[justinchuby]

Suggested change

	# producer_version='2.7.0.dev20250124+cu124',
	# producer_version='2.6.0',

[justinchuby]

Suggested change

	# producer_version='2.7.0.dev20250124+cu124',
	# producer_version='2.6.0',

[justinchuby]

I don’t see print outputs here either https://pytorch.org/tutorials/beginner/onnx/intro_onnx.html. Do we need to do something to enable it? (If there is an option)

[justinchuby]

They have prints here https://pytorch.org/tutorials/intermediate/torchrec_intro_tutorial.html

[titaiwangms]

They have prints here https://pytorch.org/tutorials/intermediate/torchrec_intro_tutorial.html

Yeah, I don't know. I tried a local build, and registry has the print out while control flow does not.

[justinchuby]

you