Add GraphBuilder.subgraph() and TensorType.to_ir() for control-flow ops

[gramalingam]

Summary:

Adds utilities to simplify building ONNX control-flow ops (Scan, Loop, If) that use graph-valued attributes.

New: TensorType.to_ir()

Converts tensor-type annotations to ir.TypeAndShape, enabling the convenient FLOAT[1024] / FLOAT['M', 'N'] /
FLOAT[...] notation wherever a type spec is needed.

New: GraphBuilder.subgraph()

Builds an ir.Graph suitable for use as a graph-valued attribute. The body is defined by a trace function fn(op,
*inputs) — the same imperative style as the outer graph. The opset version is inherited from the parent GraphBuilder
.

   def cumsum_body(op, state, x_i):
       new_state = op.Add(state, x_i)
       return new_state, new_state

   body = builder.subgraph(
       cumsum_body,
       input_types=[FLOAT[4], FLOAT[4]],
       output_types=[FLOAT[4], FLOAT[4]],
       name="cumsum_body",
   )

   final_state, partial_sums = op.Scan(
       init_state, sequence,
       body=body,
       num_scan_inputs=1,
       _outputs=2,
   )

Files changed:

• onnxscript/onnx_types.py — Add TensorType.to_ir() classmethod
• onnxscript/_internal/builder.py — Add GraphBuilder.subgraph() method + TypeSpec / _resolve_type_spec helpers
• onnxscript/onnx_types_test.py — New test file for TensorType.to_ir()
• onnxscript/_internal/builder_test.py — Tests for GraphBuilder.subgraph()
• docs/tutorial/builder/graph_builder.md — New Building Subgraphs for Control-Flow Ops section with Scan example

[codecov]

Codecov Report

❌ Patch coverage is 97.60000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.78%. Comparing base (f6a7de1) to head (53ed438).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/onnx_types_test.py	94.28%	1 Missing and 1 partial ⚠️
onnxscript/_internal/builder_test.py	98.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2824      +/-   ##
==========================================
+ Coverage   71.66%   71.78%   +0.11%     
==========================================
  Files         238      239       +1     
  Lines       28864    28989     +125     
  Branches     2849     2859      +10     
==========================================
+ Hits        20686    20809     +123     
- Misses       7207     7209       +2     
  Partials      971      971              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[justinchuby]

Nice! Thanks

[justinchuby]

Is there an option to not specify input_types and output_types? I don't think they are required for sub-graphs, right?

[gramalingam]

Is there an option to not specify input_types and output_types? I don't think they are required for sub-graphs, right?

Yes, that would be useful. I thought about it, but decided to leave it as strict for now. Will get back to it later.