Skip to content

Commit fe50b83

Browse files
gramalingamCopilot
andauthored
Add support for traced if statements in onnxscript script (#2644)
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](https://github.com/microsoft/onnxscript/blob/8a94ad646440f462dd9ae1de6b303fe5f7b7f564/onnxscript/rewriter/ort_fusions/sdpa_test.py#L40). This supports just a very basic version, where the if-condition is an outer-scope variable, like below: ```py 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). --------- Signed-off-by: Ganesan Ramalingam <grama@microsoft.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 9e0366c commit fe50b83

File tree

4 files changed

+337
-172
lines changed

4 files changed

+337
-172
lines changed

0 commit comments

Comments
 (0)