onnxscript 0.6.3.dev20260415__tar.gz → 0.6.3.dev20260417__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (197) hide show
  1. {onnxscript-0.6.3.dev20260415/onnxscript.egg-info → onnxscript-0.6.3.dev20260417}/PKG-INFO +2 -2
  2. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/builder.py +332 -117
  3. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417/onnxscript.egg-info}/PKG-INFO +2 -2
  4. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/LICENSE +0 -0
  5. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/MANIFEST.in +0 -0
  6. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/README.md +0 -0
  7. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/VERSION +0 -0
  8. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/__init__.py +0 -0
  9. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/__init__.py +0 -0
  10. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_11.py +0 -0
  11. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_5.py +0 -0
  12. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_6.py +0 -0
  13. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_7.py +0 -0
  14. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_8.py +0 -0
  15. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_framework_apis/torch_2_9.py +0 -0
  16. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/__init__.py +0 -0
  17. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/_inference.py +0 -0
  18. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/_inliner.py +0 -0
  19. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/analysis.py +0 -0
  20. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/ast_utils.py +0 -0
  21. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/autocast.py +0 -0
  22. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/converter.py +0 -0
  23. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/deprecation.py +0 -0
  24. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/evaluator.py +0 -0
  25. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/irbuilder.py +0 -0
  26. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/main.py +0 -0
  27. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/param_manipulation.py +0 -0
  28. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/sourceinfo.py +0 -0
  29. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/type_annotation.py +0 -0
  30. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/utils.py +0 -0
  31. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/values.py +0 -0
  32. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/_internal/version_utils.py +0 -0
  33. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/backend/__init__.py +0 -0
  34. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/backend/onnx_backend.py +0 -0
  35. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/backend/onnx_export.py +0 -0
  36. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/evaluator.py +0 -0
  37. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/tools/torch_lib/deduce_type_constraints.py +0 -0
  38. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/tools/torch_lib/generate_aten_signatures.py +0 -0
  39. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/tools/torch_lib/generate_prims_signatures.py +0 -0
  40. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/__init__.py +0 -0
  41. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/_constants.py +0 -0
  42. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/_flags.py +0 -0
  43. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/graph_building/__init__.py +0 -0
  44. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/__init__.py +0 -0
  45. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/common.py +0 -0
  46. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/core.py +0 -0
  47. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/fft.py +0 -0
  48. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/linalg.py +0 -0
  49. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/nested.py +0 -0
  50. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/nn.py +0 -0
  51. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/prims.py +0 -0
  52. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/quantized_decomposed.py +0 -0
  53. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/sparse.py +0 -0
  54. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/special.py +0 -0
  55. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/ops/vision.py +0 -0
  56. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/registration.py +0 -0
  57. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/function_libs/torch_lib/tensor_typing.py +0 -0
  58. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/__init__.py +0 -0
  59. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/_schemas.py +0 -0
  60. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/_tape.py +0 -0
  61. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/convenience.py +0 -0
  62. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/passes/__init__.py +0 -0
  63. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/ir/passes/common/__init__.py +0 -0
  64. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/nn/__init__.py +0 -0
  65. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/nn/_module.py +0 -0
  66. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/nn/_module_list.py +0 -0
  67. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/nn/_parameter.py +0 -0
  68. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/nn/_sequential.py +0 -0
  69. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/__init__.py +0 -0
  70. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset1.py +0 -0
  71. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset10.py +0 -0
  72. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset11.py +0 -0
  73. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset12.py +0 -0
  74. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset13.py +0 -0
  75. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset14.py +0 -0
  76. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset15.py +0 -0
  77. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset16.py +0 -0
  78. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset17.py +0 -0
  79. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset18.py +0 -0
  80. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset19.py +0 -0
  81. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset2.py +0 -0
  82. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset20.py +0 -0
  83. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset21.py +0 -0
  84. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset22.py +0 -0
  85. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset23.py +0 -0
  86. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset24.py +0 -0
  87. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset3.py +0 -0
  88. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset4.py +0 -0
  89. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset5.py +0 -0
  90. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset6.py +0 -0
  91. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset7.py +0 -0
  92. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset8.py +0 -0
  93. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset9.py +0 -0
  94. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml1.py +0 -0
  95. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml2.py +0 -0
  96. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml3.py +0 -0
  97. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml4.py +0 -0
  98. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml5.py +0 -0
  99. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/onnx_types.py +0 -0
  100. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/optimizer/__init__.py +0 -0
  101. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/optimizer/_constant_folding.py +0 -0
  102. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/optimizer/_optimizer.py +0 -0
  103. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/py.typed +0 -0
  104. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/__init__.py +0 -0
  105. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_basics.py +0 -0
  106. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_fusion_utils.py +0 -0
  107. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_ir_utils.py +0 -0
  108. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_matcher.py +0 -0
  109. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_pattern_ir.py +0 -0
  110. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/_rewrite_rule.py +0 -0
  111. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_bart_encoder.py +0 -0
  112. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_phi2lm.py +0 -0
  113. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_phi4lm.py +0 -0
  114. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_rotary_embedding_models.py +0 -0
  115. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_smollm_1.py +0 -0
  116. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_smollm_2.py +0 -0
  117. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_test_models.py +0 -0
  118. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_whisper_decoder.py +0 -0
  119. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/models/_whisper_encoder.py +0 -0
  120. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/onnx_fusions/__init__.py +0 -0
  121. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/onnx_fusions/_onnx_fusions.py +0 -0
  122. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/onnxruntime/__init__.py +0 -0
  123. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/onnxruntime/bfloat16_utils/bfloat16_converter.py +0 -0
  124. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/__init__.py +0 -0
  125. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/_core.py +0 -0
  126. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/_test_utils.py +0 -0
  127. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/attention.py +0 -0
  128. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/bias_gelu.py +0 -0
  129. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/cos_sin_cache.py +0 -0
  130. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/erfgelu.py +0 -0
  131. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py +0 -0
  132. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/gelu.py +0 -0
  133. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/gqa.py +0 -0
  134. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/gqa_packed_qkv.py +0 -0
  135. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/group_normalization_merge_silu.py +0 -0
  136. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/instance_to_group_normalization.py +0 -0
  137. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/mha.py +0 -0
  138. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/mha_bias.py +0 -0
  139. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/mha_scale.py +0 -0
  140. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/rms_normalization.py +0 -0
  141. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/rotary_embedding.py +0 -0
  142. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/sdpa.py +0 -0
  143. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/sdpa_via_mha.py +0 -0
  144. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/shape_optimization.py +0 -0
  145. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/skip_normalization.py +0 -0
  146. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/ort_fusions/softmax.py +0 -0
  147. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/pattern.py +0 -0
  148. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/__init__.py +0 -0
  149. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/__init__.py +0 -0
  150. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_basic_rules.py +0 -0
  151. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py +0 -0
  152. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py +0 -0
  153. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_collapse_slices.py +0 -0
  154. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_fuse_batchnorm.py +0 -0
  155. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_fuse_conv_affine.py +0 -0
  156. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_fuse_hardswish.py +0 -0
  157. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py +0 -0
  158. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_fuse_relus_clips.py +0 -0
  159. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_gemm_to_matmul_add.py +0 -0
  160. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_matmul_add_to_gemm.py +0 -0
  161. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_min_max_to_clip.py +0 -0
  162. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_no_op.py +0 -0
  163. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py +0 -0
  164. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_remove_expand_before_binary_op.py +0 -0
  165. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/common/_remove_optional_bias.py +0 -0
  166. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/fusion/__init__.py +0 -0
  167. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/fusion/_gqa.py +0 -0
  168. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/fusion/_layer_norm.py +0 -0
  169. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/fusion/_rms_normalization.py +0 -0
  170. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/rules/fusion/_rotary_embedding.py +0 -0
  171. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/rewriter/testing.py +0 -0
  172. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tensor.py +0 -0
  173. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/testing/__init__.py +0 -0
  174. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/__init__.py +0 -0
  175. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/memory_peak.py +0 -0
  176. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/transformers_models/__init__.py +0 -0
  177. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/transformers_models/llama.py +0 -0
  178. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/transformers_models/mistral.py +0 -0
  179. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/transformers_models/phi.py +0 -0
  180. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/tools/transformers_models/phi3.py +0 -0
  181. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/__init__.py +0 -0
  182. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/evaluation_utils.py +0 -0
  183. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/metadata_merger.py +0 -0
  184. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/replace.py +0 -0
  185. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/timing_utils.py +0 -0
  186. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/utils/utils.py +0 -0
  187. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/values.py +0 -0
  188. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/version_converter/__init__.py +0 -0
  189. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/version_converter/_c_api_utils.py +0 -0
  190. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript/version_converter/_version_converter.py +0 -0
  191. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript.egg-info/SOURCES.txt +0 -0
  192. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript.egg-info/dependency_links.txt +0 -0
  193. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript.egg-info/requires.txt +0 -0
  194. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/onnxscript.egg-info/top_level.txt +0 -0
  195. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/pyproject.toml +0 -0
  196. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/setup.cfg +0 -0
  197. {onnxscript-0.6.3.dev20260415 → onnxscript-0.6.3.dev20260417}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: onnxscript
3
- Version: 0.6.3.dev20260415
3
+ Version: 0.6.3.dev20260417
4
4
  Summary: Naturally author ONNX functions and models using a subset of Python
5
5
  Author-email: Microsoft Corporation <onnx@microsoft.com>
6
6
  License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
27
27
 
28
28
  Project-URL: Homepage, https://microsoft.github.io/onnxscript/
29
29
  Project-URL: Repository, https://github.com/microsoft/onnxscript
30
- Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/63ffecf91fdade45f27ac276f3e99ff403b4d227
30
+ Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/c6e8ec69ead8f2d4736e1103cdccebac83cdfad1
31
31
  Classifier: Development Status :: 4 - Beta
32
32
  Classifier: Environment :: Console
33
33
  Classifier: Intended Audience :: Developers
@@ -75,18 +75,65 @@ def _constant_name(
75
75
  return f"const_1d_{num}"
76
76
 
77
77
 
78
+ def lift_initializers_to_constants(graph: ir.Graph) -> None:
79
+ """Replace every initializer in *graph* with a ``Constant`` node.
80
+
81
+ ONNX ``ir.Function`` bodies do not support initializers — all values
82
+ must be produced by nodes. Call this on the function-body graph
83
+ **before** wrapping it in :class:`ir.Function` so that any constant
84
+ initializers (e.g. from Python literals promoted by
85
+ :class:`GraphBuilder`) become valid ``Constant`` nodes.
86
+
87
+ The function preserves ``ir.Value`` identity: it reuses each existing
88
+ ``ir.Value`` object as the output of the new ``Constant`` node so that
89
+ all downstream references remain valid.
90
+
91
+ Graph inputs that are *also* registered as initializers (the standard
92
+ ONNX pattern for optional inputs with default values) are skipped
93
+ because they are explicit function parameters, not embedded constants.
94
+ """
95
+ graph_input_set = {id(v) for v in graph.inputs}
96
+ to_lift: list[ir.Value] = [
97
+ v for v in list(graph.initializers.values()) if id(v) not in graph_input_set
98
+ ]
99
+ opset_version = graph.opset_imports.get("", 1)
100
+ new_nodes: list[ir.Node] = []
101
+ for value in to_lift:
102
+ tensor = value.const_value
103
+ if tensor is None:
104
+ raise ValueError(f"Initializer {value.name!r} has no const_value")
105
+ # Build a Constant node whose output is the *same* ir.Value so
106
+ # that every existing reference keeps working.
107
+ node = ir.Node(
108
+ "",
109
+ "Constant",
110
+ inputs=[],
111
+ attributes=[ir.Attr("value", ir.AttributeType.TENSOR, tensor)],
112
+ outputs=[value],
113
+ version=opset_version,
114
+ name=f"initializer_{value.name}",
115
+ )
116
+ graph.initializers.pop(value.name)
117
+ new_nodes.append(node)
118
+ # Insert all Constant nodes at the beginning of the graph.
119
+ if new_nodes:
120
+ first_existing = graph.node(0) if graph.num_nodes() > 0 else None
121
+ if first_existing is not None:
122
+ graph.insert_before(first_existing, new_nodes)
123
+ else:
124
+ for n in new_nodes:
125
+ graph.append(n)
126
+
127
+
78
128
  # Type accepted as an element of *inputs* / *outputs* by
79
129
  # :meth:`GraphBuilder.subgraph`. Can be an already-resolved
80
130
  # :class:`ir.TypeAndShape`, or a
81
131
  # :class:`~onnxscript.onnx_types.TensorType` subclass such as ``FLOAT[1024]``.
132
+ #
133
+ # .. deprecated::
134
+ # Use ``ir.Value`` with name/type/shape directly instead.
82
135
  TypeSpec = Union[ir.TypeAndShape, Any]
83
136
 
84
- # Acceptable collection forms for *inputs* / *outputs* in
85
- # :meth:`GraphBuilder.subgraph`. A :class:`Sequence` of :data:`TypeSpec`
86
- # auto-names entries (``input_0``, ``input_1``, …), while a :class:`Mapping`
87
- # from :class:`str` to :data:`TypeSpec` uses the keys as explicit names.
88
- InputOutputSpec = Union[Sequence[TypeSpec], Mapping[str, TypeSpec]]
89
-
90
137
 
91
138
  def _resolve_type_spec(spec: TypeSpec) -> ir.TypeAndShape:
92
139
  """Convert a *TypeSpec* to an :class:`ir.TypeAndShape`.
@@ -95,6 +142,9 @@ def _resolve_type_spec(spec: TypeSpec) -> ir.TypeAndShape:
95
142
  ``to_ir_type_and_shape()`` method (e.g. a
96
143
  :class:`~onnxscript.onnx_types.TensorType` subclass such as
97
144
  ``FLOAT[1024]`` or ``FLOAT['M', 'N']``).
145
+
146
+ .. deprecated::
147
+ Use :func:`make_value` or construct ``ir.Value`` directly instead.
98
148
  """
99
149
  if isinstance(spec, ir.TypeAndShape):
100
150
  return spec
@@ -112,26 +162,80 @@ def _resolve_type_spec(spec: TypeSpec) -> ir.TypeAndShape:
112
162
  )
113
163
 
114
164
 
115
- def _normalize_io_spec(
116
- spec: InputOutputSpec, default_prefix: str
117
- ) -> list[tuple[str, ir.TypeAndShape]]:
118
- """Normalize an *InputOutputSpec* into a list of ``(name, TypeAndShape)`` pairs.
165
+ def make_value(name: str, type_spec: TypeSpec | None = None) -> ir.Value:
166
+ """Create an :class:`ir.Value` from a name and optional :data:`TypeSpec`.
167
+
168
+ Similar to :func:`onnx_ir.val` but accepts a :data:`TypeSpec` (e.g.
169
+ ``FLOAT[3, 4]``) instead of separate *dtype* and *shape* arguments.
170
+
171
+ Example::
172
+
173
+ x = make_value("x", FLOAT[3, 4])
174
+ y = make_value("y") # untyped
175
+
176
+ Args:
177
+ name: The value name.
178
+ type_spec: Optional type specification. Accepts an
179
+ :class:`ir.TypeAndShape`, or a
180
+ :class:`~onnxscript.onnx_types.TensorType` subclass
181
+ (e.g. ``FLOAT[3, 4]``).
119
182
 
120
- When *spec* is a :class:`Mapping`, the keys are used as names. When it is
121
- a plain :class:`Sequence`, names are generated as
122
- ``{default_prefix}_0``, ``{default_prefix}_1``, etc.
183
+ Returns:
184
+ A fresh :class:`ir.Value` with the given name and optional type/shape.
123
185
  """
124
- if isinstance(spec, Mapping):
125
- return [(name, _resolve_type_spec(ts)) for name, ts in spec.items()]
126
- return [(f"{default_prefix}_{i}", _resolve_type_spec(ts)) for i, ts in enumerate(spec)]
186
+ if type_spec is not None:
187
+ ts = _resolve_type_spec(type_spec)
188
+ return ir.Value(name=name, type=ts.type, shape=ts.shape)
189
+ return ir.Value(name=name)
190
+
191
+
192
+ def _split_optional_inputs(
193
+ inputs: Sequence[ir.Value | None],
194
+ ) -> tuple[list[ir.Value | None], list[ir.Value]]:
195
+ """Split an input list into trace args and graph inputs.
196
+
197
+ For each ``None`` entry, a placeholder :class:`ir.Value` with a generated
198
+ name (``input_0``, ``input_1``, …) is created and added to
199
+ *graph_inputs* so that the function/graph signature declares the formal
200
+ parameter. The corresponding *trace_args* entry remains ``None`` so that
201
+ the trace function can branch with ``if x is None:``.
202
+
203
+ Returns:
204
+ A tuple of (trace_args, graph_inputs) where trace_args preserves
205
+ ``None`` holes and graph_inputs includes placeholders for absent
206
+ optional inputs.
207
+
208
+ Raises:
209
+ ValueError: If any non-None input already has a producer or is
210
+ already attached to a graph.
211
+ """
212
+ trace_args: list[ir.Value | None] = list(inputs)
213
+ graph_inputs: list[ir.Value] = []
214
+ for i, v in enumerate(trace_args):
215
+ if v is None:
216
+ # Placeholder: declared in function signature but unused in body.
217
+ graph_inputs.append(ir.Value(name=f"input_{i}"))
218
+ else:
219
+ if v.producer() is not None:
220
+ raise ValueError(
221
+ f"Input {v.name!r} already has a producer node. "
222
+ f"Pass freshly created ir.Value objects."
223
+ )
224
+ if v.graph is not None:
225
+ raise ValueError(
226
+ f"Input {v.name!r} is already attached to a graph. "
227
+ f"Pass freshly created ir.Value objects."
228
+ )
229
+ graph_inputs.append(v)
230
+ return trace_args, graph_inputs
127
231
 
128
232
 
129
233
  def build_graph(
130
234
  trace_function: Callable,
131
- inputs: InputOutputSpec,
132
- outputs: InputOutputSpec,
235
+ inputs: Sequence[ir.Value | None],
236
+ outputs: Sequence[ir.Value],
133
237
  *,
134
- opset_imports: dict[str, int] | None = None,
238
+ opset_imports: dict[str, int],
135
239
  name: str = "subgraph",
136
240
  parent: GraphBuilder | None = None,
137
241
  ) -> ir.Graph:
@@ -145,26 +249,28 @@ def build_graph(
145
249
 
146
250
  body = build_graph(
147
251
  lambda op, x, y: op.Add(x, y),
148
- inputs={"x": FLOAT[...], "y": FLOAT[...]},
149
- outputs={"sum": FLOAT[...]},
252
+ inputs=[make_value("x", FLOAT[3, 4]), make_value("y", FLOAT[3, 4])],
253
+ outputs=[make_value("sum", FLOAT[3, 4])],
254
+ opset_imports={"": 23},
150
255
  )
151
256
 
152
257
  Args:
153
258
  trace_function: A callable with signature
154
- ``(op: OpBuilder, *inputs: ir.Value) -> ir.Value | Sequence[ir.Value]``.
259
+ ``(op: OpBuilder, *inputs: ir.Value | None) -> ir.Value | Sequence[ir.Value]``.
155
260
  It is called once with freshly created placeholder inputs to record the
156
- graph topology.
157
- inputs: Types (and optionally names) for each graph input. May be a
158
- :class:`Sequence` of :data:`TypeSpec` values (names are auto-generated
159
- as ``input_0``, ``input_1``, …) **or** a :class:`Mapping` from
160
- :class:`str` names to :data:`TypeSpec` values. Each :data:`TypeSpec`
161
- can be an :class:`ir.TypeAndShape` or a
162
- :class:`~onnxscript.onnx_types.TensorType` subclass (e.g.
163
- ``FLOAT[1024]`` or ``FLOAT['M', 'N']``).
164
- outputs: Types (and optionally names) for each graph output, in the
165
- same format as *inputs*.
261
+ graph topology. ``None`` entries in *inputs* are passed through as ``None``
262
+ to support optional inputs.
263
+ inputs: A :class:`Sequence` of :class:`ir.Value` (or ``None`` for
264
+ absent optional inputs). Each ``ir.Value`` should be freshly
265
+ created with a name and optional type/shape. For ``None``
266
+ entries, placeholder values are declared as formal graph inputs,
267
+ while ``None`` is passed to *trace_function* for the
268
+ corresponding argument position.
269
+ outputs: A :class:`Sequence` of :class:`ir.Value` objects declaring
270
+ the expected outputs. After tracing, the name and type of each
271
+ declared output are applied to the corresponding returned value.
166
272
  opset_imports: Opset version map for the subgraph (e.g.
167
- ``{"": 23}``). Defaults to ``{"": 23}`` when *None*.
273
+ ``{"": 23}``).
168
274
  name: Name of the resulting :class:`ir.Graph`.
169
275
  parent: Optional parent :class:`GraphBuilder`. When provided, the
170
276
  sub-builder's ``_root`` points to the root builder of the parent,
@@ -177,42 +283,141 @@ def build_graph(
177
283
  passed directly as a graph-valued attribute (e.g. the ``body`` attribute of
178
284
  a ``Scan`` or ``Loop`` node).
179
285
  """
180
- if opset_imports is None:
181
- opset_imports = {"": 23}
182
- resolved_inputs = _normalize_io_spec(inputs, "input")
183
- resolved_outputs = _normalize_io_spec(outputs, "output")
286
+ trace_args, graph_inputs = _split_optional_inputs(inputs)
184
287
 
185
288
  subgraph = ir.Graph(
186
289
  name=name,
187
- inputs=[],
290
+ inputs=graph_inputs,
188
291
  outputs=[],
189
292
  nodes=[],
190
293
  opset_imports=opset_imports,
191
294
  )
192
295
 
193
- for input_name, ts in resolved_inputs:
194
- subgraph.inputs.append(ir.Value(name=input_name, type=ts.type, shape=ts.shape))
195
-
196
296
  sub_builder = GraphBuilder(subgraph, parent=parent)
197
297
  if parent is not None:
198
298
  sub_builder._scope_stack = list(parent._scope_stack)
199
- trace_outputs = trace_function(sub_builder.op, *subgraph.inputs)
299
+ trace_outputs = trace_function(sub_builder.op, *trace_args)
200
300
  if not isinstance(trace_outputs, Sequence):
201
301
  trace_outputs = [trace_outputs]
202
- if len(trace_outputs) != len(resolved_outputs):
302
+ if len(trace_outputs) != len(outputs):
203
303
  raise ValueError(
204
304
  f"trace_function returned {len(trace_outputs)} output(s), "
205
- f"but {len(resolved_outputs)} were declared in outputs."
305
+ f"but {len(outputs)} were declared in outputs."
206
306
  )
207
- for output, (output_name, ts) in zip(trace_outputs, resolved_outputs):
208
- output.name = output_name
209
- output.type = ts.type
210
- output.merge_shapes(ts.shape)
307
+ for returned_val, declared_val in zip(trace_outputs, outputs):
308
+ if declared_val.name:
309
+ returned_val.name = declared_val.name
310
+ if declared_val.type is not None:
311
+ if returned_val.type is not None and returned_val.type != declared_val.type:
312
+ raise ValueError(
313
+ f"Output {declared_val.name!r}: traced type "
314
+ f"{returned_val.type} conflicts with declared type "
315
+ f"{declared_val.type}."
316
+ )
317
+ returned_val.type = declared_val.type
318
+ if declared_val.shape is not None:
319
+ returned_val.merge_shapes(declared_val.shape)
211
320
 
212
321
  subgraph.outputs.extend(trace_outputs)
213
322
  return subgraph
214
323
 
215
324
 
325
+ def build_function(
326
+ trace_function: Callable,
327
+ inputs: Sequence[ir.Value | None],
328
+ *,
329
+ domain: str,
330
+ name: str,
331
+ attributes: Mapping[str, ir.Attr] | Sequence[ir.Attr] | None = None,
332
+ opset_imports: dict[str, int],
333
+ ) -> ir.Function:
334
+ """Build an :class:`ir.Function` by tracing *trace_function*.
335
+
336
+ This utility handles all boilerplate for constructing an ``ir.Function``:
337
+ graph creation, input/output wiring, initializer lifting (so that Python
338
+ literals work correctly inside function bodies), and attribute packaging.
339
+
340
+ Example::
341
+
342
+ fn = build_function(
343
+ lambda op, x, y: op.Add(x, y),
344
+ [make_value("x"), make_value("y")],
345
+ domain="com.example",
346
+ name="MyAdd",
347
+ opset_imports={"": 23},
348
+ )
349
+
350
+ Args:
351
+ trace_function: A callable with signature
352
+ ``(op: OpBuilder, *inputs: ir.Value | None) -> ir.Value | Sequence[ir.Value] | None``.
353
+ It is called once to trace the function body. Return value(s)
354
+ become function outputs. If ``None`` is returned, the function
355
+ uses whatever outputs were appended to ``graph.outputs`` by the
356
+ trace function directly.
357
+ inputs: A :class:`Sequence` of :class:`ir.Value` (or ``None`` for
358
+ absent optional inputs). ``None`` entries are represented by
359
+ placeholder formal inputs in the generated function signature,
360
+ while ``None`` is passed through to *trace_function* in the
361
+ corresponding positions so the body can branch with
362
+ ``if x is None``.
363
+ domain: Function domain (e.g. ``"com.microsoft"``).
364
+ name: Function name (e.g. ``"LinearAttention"``).
365
+ attributes: Function-level attributes. Accepts a
366
+ :class:`Mapping` from name to :class:`ir.Attr`, a
367
+ :class:`Sequence` of :class:`ir.Attr`, or ``None``.
368
+ opset_imports: Opset version map (e.g. ``{"": 23}``).
369
+
370
+ Returns:
371
+ An :class:`ir.Function` with initializers automatically lifted to
372
+ ``Constant`` nodes.
373
+ """
374
+ trace_args, graph_inputs = _split_optional_inputs(inputs)
375
+
376
+ graph = ir.Graph(
377
+ inputs=graph_inputs,
378
+ outputs=[],
379
+ nodes=[],
380
+ name=f"{name}_body",
381
+ opset_imports=opset_imports,
382
+ )
383
+
384
+ gb = GraphBuilder(graph) # No parent — function is self-contained
385
+ trace_outputs = trace_function(gb.op, *trace_args)
386
+
387
+ # Normalize outputs: either returned or appended directly, not both.
388
+ if trace_outputs is not None:
389
+ if not isinstance(trace_outputs, Sequence):
390
+ trace_outputs = [trace_outputs]
391
+ if graph.outputs:
392
+ raise ValueError(
393
+ "trace_function both returned output values and appended "
394
+ "to graph.outputs. Use one approach, not both."
395
+ )
396
+ graph.outputs.extend(trace_outputs)
397
+ elif not graph.outputs:
398
+ raise ValueError(
399
+ "trace_function returned None and did not append any outputs to graph.outputs."
400
+ )
401
+
402
+ # Lift initializers → Constant nodes (required for ir.Function bodies).
403
+ lift_initializers_to_constants(graph)
404
+
405
+ # Build attributes dict.
406
+ if attributes is None:
407
+ attr_dict: dict[str, ir.Attr] = {}
408
+ elif isinstance(attributes, Mapping):
409
+ attr_dict = dict(attributes)
410
+ else:
411
+ attr_dict = {a.name: a for a in attributes}
412
+
413
+ return ir.Function(
414
+ domain=domain,
415
+ name=name,
416
+ graph=graph,
417
+ attributes=attr_dict,
418
+ )
419
+
420
+
216
421
  class GraphBuilder:
217
422
  """Imperative builder for constructing ONNX IR graphs with automatic constant promotion, type casting, and shape inference."""
218
423
 
@@ -234,10 +439,13 @@ class GraphBuilder:
234
439
  # class_name is the qualified class name (e.g. "Gemma3DecoderLayer").
235
440
  self._scope_stack: list[tuple[str, str]] = []
236
441
 
237
- # Cache for constant initializers (scalars and sequences), keyed by (value, dtype).
238
- # This avoids creating duplicate initializers for the same constant
239
- # and allows sharing them across different layers/contexts.
240
- self._constant_cache: dict[tuple[Any, ir.DataType | None], ir.Value] = {}
442
+ # Cache for constant initializers (scalars and sequences), keyed by
443
+ # (value, dtype). Only the **root** builder owns a cache; child
444
+ # builders delegate to ``self._root`` so that all constant
445
+ # initializers live in the root graph (outer-scope initializers are
446
+ # visible to subgraphs per the ONNX spec).
447
+ if parent is None:
448
+ self._constant_cache: dict[tuple[Any, ir.DataType | None], ir.Value] = {}
241
449
 
242
450
  def opset(self, domain: str, version: int = 1) -> OpBuilder:
243
451
  """Create an OpBuilder bound to the given domain and version."""
@@ -264,7 +472,17 @@ class GraphBuilder:
264
472
  def initializer(
265
473
  self, tensor: ir.TensorProtocol, name: str | None = None, *, qualify: bool = True
266
474
  ) -> ir.Value:
267
- """Register a tensor as a graph initializer, returning the corresponding ir.Value."""
475
+ """Register a tensor as a graph initializer in the **root** graph.
476
+
477
+ Initializers created through this method are stored in the root graph
478
+ so that inner scopes (subgraphs) can reference them via ONNX's
479
+ outer-scope visibility rules. This does not apply to the ONNX
480
+ default-input pattern created via :meth:`input` with ``const_value``,
481
+ which registers an initializer on the owning graph. For function
482
+ bodies (which cannot have initializers), apply
483
+ :func:`lift_initializers_to_constants` before wrapping in
484
+ :class:`ir.Function`.
485
+ """
268
486
  if name is None:
269
487
  name = tensor.name
270
488
  if qualify:
@@ -273,7 +491,7 @@ class GraphBuilder:
273
491
  value = ir.Value(
274
492
  name=name, shape=shape, type=ir.TensorType(tensor.dtype), const_value=tensor
275
493
  )
276
- self._graph.register_initializer(value)
494
+ self._root._graph.register_initializer(value)
277
495
  return value
278
496
 
279
497
  def input(
@@ -324,6 +542,53 @@ class GraphBuilder:
324
542
  value.name = name
325
543
  self._graph.outputs.append(value)
326
544
 
545
+ def _get_or_create_constant(
546
+ self, value: VALUE_LIKE, dtype: ir.DataType | None
547
+ ) -> ir.Value:
548
+ """Materialise a constant as an initializer in the **root** graph.
549
+
550
+ Child builders delegate to the root so that all constant initializers
551
+ live in the root graph. For subgraphs this is correct because ONNX
552
+ allows inner scopes to reference outer-scope initializers. For
553
+ function bodies (which cannot reference outer initializers) callers
554
+ should apply :func:`lift_initializers_to_constants` before wrapping
555
+ the graph in :class:`ir.Function`.
556
+ """
557
+ root = self._root
558
+ if isinstance(value, (int, float, bool, str)):
559
+ if dtype is None:
560
+ dtype = _PYTHON_TYPE_TO_DTYPE.get(type(value))
561
+ cache_key = (value, dtype)
562
+ if cache_key in root._constant_cache:
563
+ return root._constant_cache[cache_key]
564
+ type_suffix = _dtype_suffix(dtype) if dtype is not None else ""
565
+ name = _constant_name(value, type_suffix, len(root._constant_cache))
566
+ tensor = ir.tensor(value, dtype=dtype, name=name)
567
+ ir_value = root.initializer(tensor, name=name, qualify=False)
568
+ root._constant_cache[cache_key] = ir_value
569
+ return ir_value
570
+ if (
571
+ isinstance(value, (list, tuple))
572
+ and value
573
+ and all(isinstance(v, type(value[0])) for v in value)
574
+ and isinstance(value[0], (int, float, bool, str))
575
+ ):
576
+ if dtype is None:
577
+ dtype = _PYTHON_TYPE_TO_DTYPE.get(type(value[0]))
578
+ cache_key = (tuple(value), dtype)
579
+ if cache_key in root._constant_cache:
580
+ return root._constant_cache[cache_key]
581
+ type_suffix = _dtype_suffix(dtype) if dtype is not None else ""
582
+ name = _constant_name(value, type_suffix, len(root._constant_cache))
583
+ tensor = ir.tensor(list(value), dtype=dtype, name=name)
584
+ ir_value = root.initializer(tensor, name=name, qualify=False)
585
+ root._constant_cache[cache_key] = ir_value
586
+ return ir_value
587
+ # For other types (TensorProtocol, numpy arrays, torch tensors, etc.),
588
+ # ir.tensor() handles the conversion.
589
+ # TODO(rama): Consider caching for other tensor values.
590
+ return self.initializer(ir.tensor(value, dtype=dtype))
591
+
327
592
  def _input_to_ir_value(
328
593
  self, value: VALUE_LIKE, like_type: ir.Value | None = None
329
594
  ) -> ir.Value | None:
@@ -343,48 +608,11 @@ class GraphBuilder:
343
608
  else None
344
609
  )
345
610
  needs_dynamic_cast = like_type is not None and dtype is None
346
- # For simple scalar/sequence constants, use a cache to avoid duplicate initializers.
347
- # These are shared across layers, so we don't qualify the name with context prefix.
348
- if isinstance(value, (int, float, bool, str)):
349
- # Normalize dtype: when None, use the default ONNX type for the
350
- # Python scalar so that (value, None) and (value, default_dtype)
351
- # share one cache entry and one initializer name.
352
- if dtype is None:
353
- dtype = _PYTHON_TYPE_TO_DTYPE.get(type(value))
354
- cache_key = (value, dtype)
355
- if cache_key in self._constant_cache:
356
- ir_value = self._constant_cache[cache_key]
357
- else:
358
- type_suffix = _dtype_suffix(dtype) if dtype is not None else ""
359
- name = _constant_name(value, type_suffix, len(self._constant_cache))
360
- tensor = ir.tensor(value, dtype=dtype, name=name)
361
- ir_value = self.initializer(tensor, name=name, qualify=False)
362
- self._constant_cache[cache_key] = ir_value
363
- elif (
364
- isinstance(value, (list, tuple))
365
- and value
366
- and all(isinstance(v, type(value[0])) for v in value)
367
- and isinstance(value[0], (int, float, bool, str))
368
- ):
369
- # Same normalization for sequences of scalars.
370
- if dtype is None:
371
- dtype = _PYTHON_TYPE_TO_DTYPE.get(type(value[0]))
372
- cache_key = (tuple(value), dtype)
373
- if cache_key in self._constant_cache:
374
- ir_value = self._constant_cache[cache_key]
375
- else:
376
- type_suffix = _dtype_suffix(dtype) if dtype is not None else ""
377
- name = _constant_name(value, type_suffix, len(self._constant_cache))
378
- tensor = ir.tensor(list(value), dtype=dtype, name=name)
379
- ir_value = self.initializer(tensor, name=name, qualify=False)
380
- self._constant_cache[cache_key] = ir_value
381
- else:
382
- # For other types (TensorProtocol, numpy arrays, torch tensors, etc.),
383
- # ir.tensor() handles the conversion.
384
- # TODO(rama): Consider caching for other tensor values.
385
- ir_value = self.initializer(ir.tensor(value, dtype=dtype))
611
+ ir_value = self._get_or_create_constant(value, dtype)
386
612
  # If like_type is provided but its type is unknown, insert a dynamic CastLike
387
613
  # so the constant is cast to match like_type's type at runtime.
614
+ # The CastLike node is created in THIS builder's graph (not root),
615
+ # so that it lives in the correct scope (subgraph or function body).
388
616
  if needs_dynamic_cast:
389
617
  ir_value = self.op.CastLike(ir_value, like_type)
390
618
  return ir_value
@@ -517,8 +745,8 @@ class GraphBuilder:
517
745
  def subgraph(
518
746
  self,
519
747
  trace_function: Callable,
520
- inputs: InputOutputSpec,
521
- outputs: InputOutputSpec,
748
+ inputs: Sequence[ir.Value | None],
749
+ outputs: Sequence[ir.Value],
522
750
  *,
523
751
  name: str = "subgraph",
524
752
  ) -> ir.Graph:
@@ -532,33 +760,20 @@ class GraphBuilder:
532
760
 
533
761
  body = graph_builder.subgraph(
534
762
  lambda op, x, y: op.Add(x, y),
535
- inputs=[FLOAT[...], FLOAT[...]],
536
- outputs=[FLOAT[...]],
537
- )
538
-
539
- Inputs and outputs can also be given as a :class:`dict` to assign
540
- explicit names::
541
-
542
- body = graph_builder.subgraph(
543
- lambda op, x, y: op.Add(x, y),
544
- inputs={"x": FLOAT[...], "y": FLOAT[...]},
545
- outputs={"sum": FLOAT[...]},
763
+ inputs=[make_value("x", FLOAT[...]), make_value("y", FLOAT[...])],
764
+ outputs=[make_value("sum", FLOAT[...])],
546
765
  )
547
766
 
548
767
  Args:
549
768
  trace_function: A callable with signature
550
- ``(op: OpBuilder, *inputs: ir.Value) -> ir.Value | Sequence[ir.Value]``.
769
+ ``(op: OpBuilder, *inputs: ir.Value | None) -> ir.Value | Sequence[ir.Value]``.
551
770
  It is called once with freshly created placeholder inputs to record the
552
771
  graph topology.
553
- inputs: Types (and optionally names) for each graph input. May be a
554
- :class:`Sequence` of :data:`TypeSpec` values (names are auto-generated
555
- as ``input_0``, ``input_1``, …) **or** a :class:`Mapping` from
556
- :class:`str` names to :data:`TypeSpec` values. Each :data:`TypeSpec`
557
- can be an :class:`ir.TypeAndShape` or a
558
- :class:`~onnxscript.onnx_types.TensorType` subclass (e.g.
559
- ``FLOAT[1024]`` or ``FLOAT['M', 'N']``).
560
- outputs: Types (and optionally names) for each graph output, in the
561
- same format as *inputs*.
772
+ inputs: A :class:`Sequence` of :class:`ir.Value` (or ``None``
773
+ for absent optional inputs). Each ``ir.Value`` should be
774
+ freshly created with a name and optional type/shape.
775
+ outputs: A :class:`Sequence` of :class:`ir.Value` objects
776
+ declaring the expected outputs.
562
777
  name: Name of the resulting :class:`ir.Graph`.
563
778
 
564
779
  Returns:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: onnxscript
3
- Version: 0.6.3.dev20260415
3
+ Version: 0.6.3.dev20260417
4
4
  Summary: Naturally author ONNX functions and models using a subset of Python
5
5
  Author-email: Microsoft Corporation <onnx@microsoft.com>
6
6
  License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
27
27
 
28
28
  Project-URL: Homepage, https://microsoft.github.io/onnxscript/
29
29
  Project-URL: Repository, https://github.com/microsoft/onnxscript
30
- Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/63ffecf91fdade45f27ac276f3e99ff403b4d227
30
+ Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/c6e8ec69ead8f2d4736e1103cdccebac83cdfad1
31
31
  Classifier: Development Status :: 4 - Beta
32
32
  Classifier: Environment :: Console
33
33
  Classifier: Intended Audience :: Developers