Add support for traced if statements in onnxscript script

[gramalingam]

Extend the onnxscript converter to handle conditionals that are evaluated at script-time (that is, in the style of trace-mode). This makes it easier to define parametric scripts that can be used to generate variations of a pattern: for example, like the many variations of the SDPA pattern test cases.

This supports just a very basic version, where the if-condition is an outer-scope variable, like below:

   if outer_scope_variable:
      ...
   else:
      ...

For such cases, the script will just include the then or else branch as appropriate, without generating an if-node.

Also: introduce an analyzer class to encapsulate analysis information, and avoid updates to AST node.

TODO: some simple extension may be useful (perhaps allow any expression in the if-condition that does not contain local variables).

[codecov]

Codecov Report

❌ Patch coverage is 82.91667% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.47%. Comparing base (8a94ad6) to head (5717723).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/_internal/analysis.py	86.84%	14 Missing and 6 partials ⚠️
onnxscript/_internal/analysis_test.py	71.79%	10 Missing and 1 partial ⚠️
onnxscript/converter_test.py	61.90%	8 Missing ⚠️
onnxscript/converter.py	92.85%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2644      +/-   ##
==========================================
+ Coverage   70.46%   70.47%   +0.01%     
==========================================
  Files         224      224              
  Lines       26572    26684     +112     
  Branches     2637     2663      +26     
==========================================
+ Hits        18723    18806      +83     
- Misses       6928     6953      +25     
- Partials      921      925       +4     

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

[Copilot]

Pull Request Overview

This PR extends the onnxscript converter to support traced if-statements, where the condition is evaluated at script-definition time rather than runtime. This enables parametric script generation for creating pattern variations. The implementation introduces an AstAnalyzer class to encapsulate analysis information and avoid mutating AST nodes directly.

Key Changes:

• Introduces AstAnalyzer class to centralize AST analysis and store analysis results in dictionaries instead of AST node attributes
• Adds support for constant if-condition detection when the condition is an outer-scope variable
• Updates if-statement translation to emit only the relevant branch when the condition is constant

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
onnxscript/converter.py	Integrates AstAnalyzer into the converter and implements traced if-statement handling
onnxscript/_internal/analysis.py	Implements AstAnalyzer class with constant if-condition detection and refactored analysis methods
onnxscript/_internal/analysis_test.py	Updates tests to use AstAnalyzer and adds test coverage for constant if-conditions
onnxscript/converter_test.py	Adds integration test for traced if-statement functionality

[Copilot]

The type: ignore comment should include a specific error code (e.g., type: ignore[arg-type]) to make it clear what type checking issue is being suppressed.

Suggested change

	self.results.append(live_out) # type: ignore
	self.results.append(live_out) # type: ignore[assignment]

[gramalingam]

One design question for consideration here: when should we consider "tracing" the if-then-else (emitting only the then/else branch instead of an onnx If op node). Obviously, we need the condition to be evaluable at script time. The choice is whether we should require users to explicitly indicate this or whether we should do this implicitly.

In the explicit mode, we would require users to indicate via some syntax like if onnxscript.eval(cond1): to indicate that cond1 should be evaluated at script time for a tracing mode translation. In the implicit mode, we just see if the condition can be evaluated or not.

While the difference is not as significant here (implicit mode seems fine), it could be more important when we get to loops. There the question is whether the loop should be unrolled or not ... but even if a loop can be unrolled, a user might not want the loop to be unrolled, and so may prefer to exercise control over the behavior. In a loop, we might support a syntax like for var in onnxscript.unroll(iterator):.

[gramalingam]

I guess we could make it a user-controllable option to the converter with a particular default value. But it would still be useful to find out what users prefer as the default setting.

[titaiwangms]

Just for my education, is this design question for future reference? Or the current PR is implementing it in an implicit way? (Looks like right now we only consider the condition being outer-scope variable.)

I remember in torchscript that users use torch.jit.script() to specify the scripted blocks. But in general, tracing is preferred to scripting. What is the case that users might not want the loop being unrolled even it's fixed?