onnxscript 0.1.0.dev20231127__tar.gz → 0.1.0.dev20231129__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 (201) hide show
  1. {onnxscript-0.1.0.dev20231127/onnxscript.egg-info → onnxscript-0.1.0.dev20231129}/PKG-INFO +1 -1
  2. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/onnx_export.py +12 -19
  3. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/_flags.py +4 -0
  4. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/graph_building.py +1 -0
  5. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/core.py +36 -33
  6. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/fft.py +12 -5
  7. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/nn.py +10 -2
  8. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/registration.py +14 -0
  9. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +74 -12
  10. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +51 -8
  11. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/values.py +3 -0
  12. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129/onnxscript.egg-info}/PKG-INFO +1 -1
  13. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/LICENSE +0 -0
  14. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/MANIFEST.in +0 -0
  15. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/README.md +0 -0
  16. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/VERSION +0 -0
  17. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/__init__.py +0 -0
  18. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/__init__.py +0 -0
  19. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/analysis.py +0 -0
  20. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/analysis_test.py +0 -0
  21. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/ast_utils.py +0 -0
  22. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/autocast.py +0 -0
  23. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/deprecation.py +0 -0
  24. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/param_manipulation.py +0 -0
  25. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/param_manipulation_test.py +0 -0
  26. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/runtime_typing.py +0 -0
  27. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/utils.py +0 -0
  28. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/_internal/version_utils.py +0 -0
  29. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/__init__.py +0 -0
  30. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/onnx_backend.py +0 -0
  31. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/onnx_backend_test.py +0 -0
  32. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/onnx_export_test.py +0 -0
  33. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/converter.py +0 -0
  34. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/converter_test.py +0 -0
  35. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/__init__.py +0 -0
  36. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/_infra.py +0 -0
  37. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/context.py +0 -0
  38. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/decorator.py +0 -0
  39. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/formatter.py +0 -0
  40. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/__init__.py +0 -0
  41. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_address.py +0 -0
  42. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_artifact.py +0 -0
  43. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_artifact_change.py +0 -0
  44. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_artifact_content.py +0 -0
  45. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_artifact_location.py +0 -0
  46. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_attachment.py +0 -0
  47. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_code_flow.py +0 -0
  48. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_configuration_override.py +0 -0
  49. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_conversion.py +0 -0
  50. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_edge.py +0 -0
  51. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_edge_traversal.py +0 -0
  52. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_exception.py +0 -0
  53. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_external_properties.py +0 -0
  54. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_external_property_file_reference.py +0 -0
  55. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_external_property_file_references.py +0 -0
  56. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_fix.py +0 -0
  57. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_graph.py +0 -0
  58. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_graph_traversal.py +0 -0
  59. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_invocation.py +0 -0
  60. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_location.py +0 -0
  61. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_location_relationship.py +0 -0
  62. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_logical_location.py +0 -0
  63. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_message.py +0 -0
  64. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_multiformat_message_string.py +0 -0
  65. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_node.py +0 -0
  66. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_notification.py +0 -0
  67. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_physical_location.py +0 -0
  68. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_property_bag.py +0 -0
  69. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_rectangle.py +0 -0
  70. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_region.py +0 -0
  71. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_replacement.py +0 -0
  72. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_reporting_configuration.py +0 -0
  73. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_reporting_descriptor.py +0 -0
  74. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_reporting_descriptor_reference.py +0 -0
  75. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_reporting_descriptor_relationship.py +0 -0
  76. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_result.py +0 -0
  77. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_result_provenance.py +0 -0
  78. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_run.py +0 -0
  79. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_run_automation_details.py +0 -0
  80. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_sarif_log.py +0 -0
  81. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_special_locations.py +0 -0
  82. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_stack.py +0 -0
  83. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_stack_frame.py +0 -0
  84. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_suppression.py +0 -0
  85. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_thread_flow.py +0 -0
  86. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_thread_flow_location.py +0 -0
  87. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_tool.py +0 -0
  88. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_tool_component.py +0 -0
  89. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_tool_component_reference.py +0 -0
  90. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_translation_metadata.py +0 -0
  91. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_version_control_details.py +0 -0
  92. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_web_request.py +0 -0
  93. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/_web_response.py +0 -0
  94. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/sarif/version.py +0 -0
  95. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/diagnostics/infra/utils.py +0 -0
  96. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/evaluator.py +0 -0
  97. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/evaluator_test.py +0 -0
  98. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/tools/torch_lib/deduce_type_constraints.py +0 -0
  99. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py +0 -0
  100. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/tools/torch_lib/generate_aten_signatures.py +0 -0
  101. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/tools/torch_lib/generate_prims_signatures.py +0 -0
  102. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/__init__.py +0 -0
  103. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/_constants.py +0 -0
  104. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/graph_building_test.py +0 -0
  105. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/__init__.py +0 -0
  106. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/common.py +0 -0
  107. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/linalg.py +0 -0
  108. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/nested.py +0 -0
  109. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/prims.py +0 -0
  110. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/sparse.py +0 -0
  111. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/special.py +0 -0
  112. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/tensor_typing.py +0 -0
  113. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/irbuilder.py +0 -0
  114. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/main.py +0 -0
  115. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/__init__.py +0 -0
  116. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset1.py +0 -0
  117. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset10.py +0 -0
  118. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset11.py +0 -0
  119. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset12.py +0 -0
  120. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset13.py +0 -0
  121. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset14.py +0 -0
  122. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset15.py +0 -0
  123. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset16.py +0 -0
  124. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset17.py +0 -0
  125. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset18.py +0 -0
  126. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset19.py +0 -0
  127. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset2.py +0 -0
  128. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset3.py +0 -0
  129. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset4.py +0 -0
  130. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset5.py +0 -0
  131. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset6.py +0 -0
  132. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset7.py +0 -0
  133. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset8.py +0 -0
  134. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset9.py +0 -0
  135. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml1.py +0 -0
  136. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml2.py +0 -0
  137. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset_ai_onnx_ml3.py +0 -0
  138. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_opset/_impl/opset_ai_onnx_preview_training1.py +0 -0
  139. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/onnx_types.py +0 -0
  140. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/py.typed +0 -0
  141. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/sourceinfo.py +0 -0
  142. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tensor.py +0 -0
  143. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tensor_test.py +0 -0
  144. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/testing.py +0 -0
  145. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/__init__.py +0 -0
  146. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/common/__init__.py +0 -0
  147. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/common/onnx_script_test_case.py +0 -0
  148. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/common/testutils.py +0 -0
  149. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/eager_mode_test.py +0 -0
  150. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/eager_test.py +0 -0
  151. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/external_tensor_test.py +0 -0
  152. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/error_reproduction.py +0 -0
  153. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/ops_test.py +0 -0
  154. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/ops_test_common.py +0 -0
  155. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/attr_test.py +0 -0
  156. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/gemmgelu.py +0 -0
  157. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/gemmgelu_test.py +0 -0
  158. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/if_test.py +0 -0
  159. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/onnxfns1A_test.py +0 -0
  160. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/onnxfns2_test.py +0 -0
  161. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/onnxfns_test.py +0 -0
  162. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/functions/ort_custom_ops.py +0 -0
  163. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/if_test.py +0 -0
  164. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/loop_test.py +0 -0
  165. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/__init__.py +0 -0
  166. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/attrref.py +0 -0
  167. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/cast_like.py +0 -0
  168. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/different_opset.py +0 -0
  169. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/dropout.py +0 -0
  170. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/eager_op.py +0 -0
  171. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/eg1.py +0 -0
  172. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/getitem.py +0 -0
  173. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/graph_attr.py +0 -0
  174. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/identity.py +0 -0
  175. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/if_statement.py +0 -0
  176. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/loops_break.py +0 -0
  177. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/loops_while.py +0 -0
  178. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/m1.py +0 -0
  179. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/multi.py +0 -0
  180. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/onnxfns1.py +0 -0
  181. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/onnxfns1A.py +0 -0
  182. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/onnxfns2.py +0 -0
  183. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/opt_input.py +0 -0
  184. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/opt_output.py +0 -0
  185. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/renaming.py +0 -0
  186. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/sequences.py +0 -0
  187. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/signal_dft.py +0 -0
  188. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/subfunction.py +0 -0
  189. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/models/type_double.py +0 -0
  190. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/onnx_types_test.py +0 -0
  191. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/operator_test.py +0 -0
  192. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/type_annotation.py +0 -0
  193. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/type_annotation_test.py +0 -0
  194. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript/values_test.py +0 -0
  195. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript.egg-info/SOURCES.txt +0 -0
  196. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript.egg-info/dependency_links.txt +0 -0
  197. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript.egg-info/requires.txt +0 -0
  198. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/onnxscript.egg-info/top_level.txt +0 -0
  199. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/pyproject.toml +0 -0
  200. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/setup.cfg +0 -0
  201. {onnxscript-0.1.0.dev20231127 → onnxscript-0.1.0.dev20231129}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: onnxscript
3
- Version: 0.1.0.dev20231127
3
+ Version: 0.1.0.dev20231129
4
4
  Summary: Naturally author ONNX functions and models using a subset of Python
5
5
  Home-page: https://onnxscript.ai/
6
6
  Author-email: Microsoft Corporation <onnx@microsoft.com>
@@ -258,9 +258,9 @@ class Exporter:
258
258
  return str(self._rename_variable(name))
259
259
 
260
260
  def _rename_domain(self, domain: str) -> str:
261
- if domain == "":
262
- return "opset"
263
- return domain.replace(".", "_")
261
+ if domain in {"", "ai.onnx"}:
262
+ return "opset" # TODO: Need checks to avoid name conflicts.
263
+ return _cleanup_variable_name(domain) # type: ignore[return-value]
264
264
 
265
265
  def _make_opset_name(self, domain, version):
266
266
  return f"{self._rename_domain(domain)}{version}"
@@ -552,11 +552,13 @@ class Exporter:
552
552
  add_line(f" return {return_values}")
553
553
  return "\n".join(result)
554
554
 
555
- def _translate_graph(self, model: onnx.ModelProto, function_name: str) -> str:
555
+ def _translate_graph(self, model: onnx.ModelProto, function_name: Optional[str]) -> str:
556
556
  graph = model.graph
557
557
  opsets = {}
558
558
  for imported in model.opset_import:
559
559
  opsets[imported.domain] = imported.version
560
+ if function_name is None:
561
+ function_name = _cleanup_variable_name(graph.name)
560
562
 
561
563
  result: list[str] = []
562
564
 
@@ -593,7 +595,9 @@ class Exporter:
593
595
  return "from onnxscript.onnx_types import " + ", ".join(sorted_types)
594
596
  return ""
595
597
 
596
- def export(self, proto: onnx.ModelProto | onnx.FunctionProto, function_name: str) -> str:
598
+ def export(
599
+ self, proto: onnx.ModelProto | onnx.FunctionProto, function_name: Optional[str]
600
+ ) -> str:
597
601
  result: list[str] = []
598
602
 
599
603
  def add(line: str) -> None:
@@ -612,7 +616,6 @@ class Exporter:
612
616
  translated_functions.append(self._translate_graph(proto, function_name))
613
617
  else:
614
618
  assert isinstance(proto, FunctionProto)
615
- # TODO: use function_name?
616
619
  translated_functions = [self._translate_function(proto)]
617
620
 
618
621
  # TODO: unique_function_domain_version.add((f.domain, 1))
@@ -655,22 +658,15 @@ def _attribute_param_types(
655
658
 
656
659
  def export2python(
657
660
  model_onnx,
658
- opset=None,
659
- verbose=True,
660
- name=None,
661
- rename=False,
662
- function_name="main",
663
- use_operators=False,
661
+ function_name: Optional[str] = None,
662
+ rename: bool = False,
663
+ use_operators: bool = False,
664
664
  inline_const: bool = False,
665
665
  ):
666
666
  """Exports an ONNX model to the *python* syntax.
667
667
 
668
668
  Args:
669
669
  model_onnx: string or ONNX graph
670
- opset: opset to export to (None to select the one from the
671
- graph)
672
- verbose: inserts prints
673
- name: to overwrite onnx name
674
670
  rename: rename the names to get shorter names
675
671
  function_name: main function name
676
672
  use_operators: use Python operators.
@@ -694,9 +690,6 @@ def export2python(
694
690
  code = export2python(onx)
695
691
  print(code)
696
692
  """
697
- del opset # unused
698
- del verbose # unused
699
- del name # unused
700
693
  if isinstance(model_onnx, str):
701
694
  model_onnx = onnx.load(model_onnx)
702
695
 
@@ -39,3 +39,7 @@ EXPERIMENTAL_INITIALIZERS_AS_INPUTS: bool = _load_boolean_flag(
39
39
  "TORCHLIB_EXPERIMENTAL_INITIALIZERS_AS_INPUTS",
40
40
  this_will="make initializers as inputs to the model graph",
41
41
  )
42
+ EXPERIMENTAL_PREFER_TRACING: bool = _load_boolean_flag(
43
+ "TORCHLIB_EXPERIMENTAL_PREFER_TRACING",
44
+ this_will="trace all traceable functions to fold if branches and collapse constant expressions",
45
+ )
@@ -816,6 +816,7 @@ class TorchScriptGraph:
816
816
  continue
817
817
  if prefix:
818
818
  name = f"{prefix}/{name}"
819
+ value_info.name = name
819
820
  named_value_info[name] = value_info
820
821
  for name, sub_graph in self._sub_torch_script_graphs.items():
821
822
  named_value_info.update(
@@ -86,7 +86,7 @@ def aten__log_softmax_half(
86
86
  return aten__log_softmax(self, dim, half_to_float)
87
87
 
88
88
 
89
- @torch_op("aten::_log_softmax")
89
+ @torch_op("aten::_log_softmax", traceable=True)
90
90
  def aten__log_softmax(
91
91
  self: TFloatHighPrecision,
92
92
  dim: int,
@@ -251,7 +251,7 @@ def aten_addmv(
251
251
  return op.Add(op.Mul(self, beta), op.Mul(op.MatMul(mat, vec), alpha))
252
252
 
253
253
 
254
- @torch_op("aten::addr")
254
+ @torch_op("aten::addr", traceable=True)
255
255
  def aten_addr(
256
256
  self: TReal, vec1: TReal, vec2: TReal, beta: float = 1.0, alpha: float = 1.0
257
257
  ) -> TReal:
@@ -329,7 +329,7 @@ def aten_align_to(self: TensorType, names: Sequence[str]) -> TensorType:
329
329
  raise NotImplementedError()
330
330
 
331
331
 
332
- @torch_op("aten::all")
332
+ @torch_op("aten::all", traceable=True)
333
333
  def aten_all(self: TTensor) -> BOOL:
334
334
  """all(Tensor self) -> Tensor"""
335
335
 
@@ -343,7 +343,7 @@ def aten_all(self: TTensor) -> BOOL:
343
343
  return result
344
344
 
345
345
 
346
- @torch_op("aten::all.dim")
346
+ @torch_op("aten::all.dim", traceable=True)
347
347
  def aten_all_dim(self: TTensor, dim: int, keepdim: bool = False) -> BOOL:
348
348
  """all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor"""
349
349
 
@@ -371,7 +371,7 @@ def aten_all_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
371
371
  return self
372
372
 
373
373
 
374
- @torch_op("aten::all.dims")
374
+ @torch_op("aten::all.dims", traceable=True)
375
375
  def aten_all_dims_no_dim(self: TTensor, keepdims: bool) -> BOOL:
376
376
  """all.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor"""
377
377
 
@@ -451,7 +451,7 @@ def aten_angle(self: TensorType) -> TensorType:
451
451
  raise NotImplementedError()
452
452
 
453
453
 
454
- @torch_op("aten::any")
454
+ @torch_op("aten::any", traceable=True)
455
455
  def aten_any(self: TTensor) -> BOOL:
456
456
  """any(Tensor self) -> Tensor"""
457
457
 
@@ -466,7 +466,7 @@ def aten_any(self: TTensor) -> BOOL:
466
466
  return result
467
467
 
468
468
 
469
- @torch_op("aten::any.dim")
469
+ @torch_op("aten::any.dim", traceable=True)
470
470
  def aten_any_dim(self: TTensor, dim: int, keepdim: bool = False) -> BOOL:
471
471
  """any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor"""
472
472
 
@@ -496,7 +496,7 @@ def aten_any_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
496
496
  return self
497
497
 
498
498
 
499
- @torch_op("aten::any.dims")
499
+ @torch_op("aten::any.dims", traceable=True)
500
500
  def aten_any_dims_no_dim(self: TTensor, keepdims: bool) -> BOOL:
501
501
  """any.dims(Tensor self, int[1]? dim=None, bool keepdim=False) -> Tensor"""
502
502
 
@@ -688,7 +688,7 @@ def aten_arctanh(self: TensorType) -> TensorType:
688
688
  raise NotImplementedError()
689
689
 
690
690
 
691
- @torch_op("aten::argmax")
691
+ @torch_op("aten::argmax", traceable=True)
692
692
  def aten_argmax(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
693
693
  """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
694
694
 
@@ -701,7 +701,7 @@ def aten_argmax(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
701
701
  return result
702
702
 
703
703
 
704
- @torch_op("aten::argmax")
704
+ @torch_op("aten::argmax", traceable=True)
705
705
  def aten_argmax_dim(self: Union[RealType, UINT8], dim: int, keepdim: bool = False) -> INT64:
706
706
  """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
707
707
 
@@ -888,7 +888,7 @@ def aten_atanh(self: TFloat) -> TFloat:
888
888
  return op.Atanh(self)
889
889
 
890
890
 
891
- @torch_op("aten::atleast_1d")
891
+ @torch_op("aten::atleast_1d", traceable=True)
892
892
  def aten_atleast_1d(self: TTensor) -> TTensor:
893
893
  """atleast_1d(Tensor self) -> Tensor"""
894
894
 
@@ -936,7 +936,7 @@ def aten_atleast_2d_sequence(self: Sequence[TTensor]) -> TTensor:
936
936
  return op.SequenceMap(self, body=reshape_to_2d)
937
937
 
938
938
 
939
- @torch_op("aten::atleast_3d")
939
+ @torch_op("aten::atleast_3d", traceable=True)
940
940
  def aten_atleast_3d(self: TTensor) -> TTensor:
941
941
  """atleast_3d(Tensor self) -> Tensor"""
942
942
 
@@ -1548,7 +1548,7 @@ def aten_clamp(self: TReal, min: Optional[TReal] = None, max: Optional[TReal] =
1548
1548
  return clamped
1549
1549
 
1550
1550
 
1551
- @torch_op("aten::clamp_max")
1551
+ @torch_op("aten::clamp_max", traceable=True)
1552
1552
  def aten_clamp_max(self: TReal, max_: TReal) -> TReal:
1553
1553
  """clamp_max(Tensor self, Tensor max) -> Tensor"""
1554
1554
 
@@ -1567,7 +1567,7 @@ def aten_clamp_max(self: TReal, max_: TReal) -> TReal:
1567
1567
  return result
1568
1568
 
1569
1569
 
1570
- @torch_op("aten::clamp_min")
1570
+ @torch_op("aten::clamp_min", traceable=True)
1571
1571
  def aten_clamp_min(self: TReal, min_: TReal) -> TReal:
1572
1572
  """clamp_min(Tensor self, Tensor min) -> Tensor"""
1573
1573
 
@@ -1946,7 +1946,7 @@ def aten_convolution(
1946
1946
  return result
1947
1947
 
1948
1948
 
1949
- @torch_op("aten::convolution", private=True)
1949
+ @torch_op("aten::convolution", private=True, traceable=True)
1950
1950
  def _aten_convolution_onnx(
1951
1951
  input: TFloat,
1952
1952
  weight: TFloat,
@@ -2366,7 +2366,7 @@ def aten_cumsum(
2366
2366
  return _aten_cumsum_onnx(cast, dim)
2367
2367
 
2368
2368
 
2369
- @torch_op("aten::cumsum", private=True)
2369
+ @torch_op("aten::cumsum", private=True, traceable=True)
2370
2370
  def _aten_cumsum_onnx(
2371
2371
  self: TRealUnlessInt16OrInt8, dim: Union[INT32, INT64]
2372
2372
  ) -> TRealUnlessInt16OrInt8:
@@ -2451,7 +2451,7 @@ def aten_diagonal(self: TReal, offset: int = 0, dim1: int = 0, dim2: int = 1) ->
2451
2451
  return _aten_diagonal_onnx(self, offset, dim1, dim2, perm, axes)
2452
2452
 
2453
2453
 
2454
- @torch_op("aten::diagonal", private=True)
2454
+ @torch_op("aten::diagonal", private=True, traceable=True)
2455
2455
  def _aten_diagonal_onnx(
2456
2456
  self: TTensor, offset: int, dim1: int, dim2: int, perm: Sequence[int], axes: Sequence[int]
2457
2457
  ) -> TTensor:
@@ -2727,7 +2727,7 @@ def aten_dot(self: TFloat, tensor: TFloat) -> TFloat:
2727
2727
  return op.MatMul(self, tensor)
2728
2728
 
2729
2729
 
2730
- @torch_op("aten::dropout")
2730
+ @torch_op("aten::dropout", traceable=True)
2731
2731
  def aten_dropout(input: TFloat, p: FLOAT, train: BOOL) -> TFloat:
2732
2732
  """dropout(Tensor input, float p, bool train) -> Tensor"""
2733
2733
 
@@ -3062,7 +3062,7 @@ def aten_embedding_dense_backward(
3062
3062
  raise NotImplementedError()
3063
3063
 
3064
3064
 
3065
- @torch_op("aten::embedding_renorm")
3065
+ @torch_op("aten::embedding_renorm", traceable=True)
3066
3066
  def aten_embedding_renorm(
3067
3067
  weight: TFloat, indices: INT64, max_norm: float, norm_type: float = 2.0
3068
3068
  ) -> TFloat:
@@ -3541,7 +3541,7 @@ def aten_fused_moving_avg_obs_fake_quant(
3541
3541
  raise NotImplementedError()
3542
3542
 
3543
3543
 
3544
- @torch_op("aten::gather")
3544
+ @torch_op("aten::gather", traceable=True)
3545
3545
  def aten_gather(
3546
3546
  self: TReal,
3547
3547
  dim: int,
@@ -4098,7 +4098,7 @@ def aten_index_reduce(
4098
4098
  raise NotImplementedError()
4099
4099
 
4100
4100
 
4101
- @torch_op("aten::index_select")
4101
+ @torch_op("aten::index_select", traceable=True)
4102
4102
  def aten_index_select(self: TTensor, dim: int, index: IntType) -> TTensor:
4103
4103
  """index_select(Tensor self, int dim, Tensor index) -> Tensor"""
4104
4104
 
@@ -4576,12 +4576,13 @@ def aten_logaddexp(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrB
4576
4576
  @torch_op("aten::logaddexp2")
4577
4577
  def aten_logaddexp2(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrBFloat16:
4578
4578
  """logaddexp2(Tensor self, Tensor other) -> Tensor"""
4579
- summation = op.Add(op.Pow(2.0, self), op.Pow(2.0, other))
4579
+ two = op.CastLike(2.0, self)
4580
+ summation = op.Add(op.Pow(two, self), op.Pow(two, other))
4580
4581
 
4581
- return op.Div(op.Log(summation), op.CastLike(op.Log(2.0), self))
4582
+ return op.Div(op.Log(summation), op.Log(two))
4582
4583
 
4583
4584
 
4584
- @torch_op("aten::logcumsumexp")
4585
+ @torch_op("aten::logcumsumexp", traceable=True)
4585
4586
  def aten_logcumsumexp(self: TFloatOrBFloat16, dim: int) -> TFloatOrBFloat16:
4586
4587
  """logcumsumexp(Tensor self, int dim) -> Tensor"""
4587
4588
 
@@ -4673,10 +4674,12 @@ def _aten_logit_onnx(self: TFloatOrBFloat16) -> TFloatOrBFloat16:
4673
4674
 
4674
4675
  @torch_op("aten::logit", private=True)
4675
4676
  def _aten_logit_clamp_onnx(self: TFloatOrBFloat16, eps: float) -> TFloatOrBFloat16:
4676
- temporary_self = op.Where(self <= 1.0 - eps, self, 1.0 - eps)
4677
+ eps = op.CastLike(eps, self)
4678
+ one = op.CastLike(1.0, self)
4679
+ temporary_self = op.Where(self <= one - eps, self, one - eps)
4677
4680
  z = op.Where(temporary_self < eps, eps, temporary_self)
4678
4681
 
4679
- return op.Log(op.Div(z, op.Sub(1.0, z)))
4682
+ return op.Log(op.Div(z, op.Sub(one, z)))
4680
4683
 
4681
4684
 
4682
4685
  @torch_op("aten::logit", trace_only=True)
@@ -4693,7 +4696,7 @@ def aten_logspace(start: float, end: float, steps: int, base: float = 10.0) -> T
4693
4696
  raise NotImplementedError()
4694
4697
 
4695
4698
 
4696
- @torch_op("aten::logsumexp")
4699
+ @torch_op("aten::logsumexp", traceable=True)
4697
4700
  def aten_logsumexp(self: TFloat, dim: INT64, keepdim: int = False) -> TFloat:
4698
4701
  """logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor"""
4699
4702
 
@@ -4915,7 +4918,7 @@ def aten_max(self: TReal) -> TReal:
4915
4918
  return result
4916
4919
 
4917
4920
 
4918
- @torch_op("aten::max.dim")
4921
+ @torch_op("aten::max.dim", traceable=True)
4919
4922
  def aten_max_dim(self: TReal, dim: int, keepdim: bool = False) -> Tuple[TReal, INT64]:
4920
4923
  """max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)"""
4921
4924
 
@@ -4951,7 +4954,7 @@ def aten_mean(self: TReal) -> TReal:
4951
4954
  return op.Squeeze(result)
4952
4955
 
4953
4956
 
4954
- @torch_op("aten::mean.dim")
4957
+ @torch_op("aten::mean.dim", traceable=True)
4955
4958
  def aten_mean_dim(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
4956
4959
  """mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor"""
4957
4960
 
@@ -4983,7 +4986,7 @@ def aten_min(self: TReal) -> TReal:
4983
4986
  return op.ReduceMin(self, keepdims=False)
4984
4987
 
4985
4988
 
4986
- @torch_op("aten::min.dim")
4989
+ @torch_op("aten::min.dim", traceable=True)
4987
4990
  def aten_min_dim(self: TReal, dim: int, keepdim: bool = False) -> Tuple[TReal, TInt]:
4988
4991
  """min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)"""
4989
4992
  if IsScalar(self):
@@ -5476,7 +5479,7 @@ def aten_nansum(
5476
5479
  raise NotImplementedError()
5477
5480
 
5478
5481
 
5479
- @torch_op("aten::narrow")
5482
+ @torch_op("aten::narrow", traceable=True)
5480
5483
  def aten_narrow(self: TTensor, dim: INT64, start: INT64, length: INT64) -> TTensor:
5481
5484
  """narrow(Tensor(a) self, int dim, SymInt start, SymInt length) -> Tensor(a)"""
5482
5485
 
@@ -6103,7 +6106,7 @@ def aten_norm_except_dim(v: TensorType, pow: int = 2, dim: int = 0) -> TensorTyp
6103
6106
  raise NotImplementedError()
6104
6107
 
6105
6108
 
6106
- @torch_op(("aten::normal", "aten::normal_functional"))
6109
+ @torch_op(("aten::normal", "aten::normal_functional"), traceable=True)
6107
6110
  def aten_normal(
6108
6111
  self: TTensor,
6109
6112
  mean: float = 0.0,
@@ -7672,7 +7675,7 @@ def aten_sum_dim_IntList(
7672
7675
  return result
7673
7676
 
7674
7677
 
7675
- @torch_op("aten::sum", private=True)
7678
+ @torch_op("aten::sum", private=True, traceable=True)
7676
7679
  def _aten_sum_dim_onnx(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
7677
7680
  self_is_scalar = IsScalar(self)
7678
7681
  if self_is_scalar:
@@ -95,11 +95,18 @@ def _fftn_onnx(
95
95
  # dimension at the beginning to represent the batch dimension.
96
96
  transformed = op.Unsqueeze(self, axes=[0])
97
97
 
98
- for dim_ in dims:
99
- if dim_ >= 0:
100
- # Add 1 to account for the batch dimension when counting axes from the left
101
- dim_ = dim_ + 1
102
- transformed = op.DFT(transformed, axis=dim_, inverse=inverse, onesided=onesided)
98
+ # Add 1 to account for the batch dimension when counting axes from the left
99
+ new_dims = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dims]
100
+
101
+ for dim in new_dims[:-1]:
102
+ transformed = op.DFT(transformed, axis=dim, inverse=inverse, onesided=False)
103
+
104
+ # Torch computers one-sided FFT on the last dimension only.
105
+ if onesided:
106
+ transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=True)
107
+ else:
108
+ transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=False)
109
+
103
110
  # Remove the batch dimension
104
111
  transformed = op.Squeeze(transformed, axes=[0])
105
112
 
@@ -1404,10 +1404,18 @@ def aten_pad_sequence(
1404
1404
  raise NotImplementedError()
1405
1405
 
1406
1406
 
1407
- def aten_reflection_pad1d(self: TensorType, padding: INT64) -> TensorType:
1407
+ @torch_op("aten::reflection_pad1d")
1408
+ def aten_reflection_pad1d(self: TFloat, padding: INT64) -> TFloat:
1408
1409
  """reflection_pad1d(Tensor self, SymInt[2] padding) -> Tensor"""
1409
1410
 
1410
- raise NotImplementedError()
1411
+ # assert len(padding) == 2
1412
+ # Input of padding argument should be [x,y], need change to onnx format [0, x, 0, y]
1413
+ start = op.Slice(padding, [0], [1], axes=[0])
1414
+ end = op.Slice(padding, [1], [2], axes=[0])
1415
+ padding_onnx = op.Concat(
1416
+ op.Constant(value_ints=[0]), start, op.Constant(value_ints=[0]), end, axis=0
1417
+ )
1418
+ return op.Pad(self, padding_onnx, mode="reflect")
1411
1419
 
1412
1420
 
1413
1421
  def aten_reflection_pad1d_backward(
@@ -99,6 +99,7 @@ def torch_op(
99
99
  trace_only: bool = False,
100
100
  private: bool = False,
101
101
  complex: bool = False,
102
+ traceable: bool = False,
102
103
  ) -> Callable[[FunctionType], onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction]:
103
104
  """Register a torch op.
104
105
 
@@ -112,6 +113,18 @@ def torch_op(
112
113
  private: Whether the function is private (not directly exposed). It should
113
114
  be true for all functions with names starting with "_".
114
115
  complex: Whether the function expects complex-valued inputs.
116
+ traceable: Whether the function can also be traced. This is an **experimental** flag.
117
+ A function is traceable if it can both be scripted and traced to produce
118
+ the same result for a given input. Specifically:
119
+
120
+ - A function _can_ be tagged with traceable if its if branches (if any)
121
+ can be statically evaluated.
122
+ - A function _should_ be tagged with traceable if it contains if branches
123
+ and/or CastLike nodes so that they can be evaluated away with the
124
+ EXPERIMENTAL_PREFER_TRACING on.
125
+ - A function without if branches or CastLike nodes _should not_ be tagged
126
+ with traceable because inlining will do the same thing.
127
+ - A function with `@graph` defined for a `Scan` op is not traceable yet.
115
128
  """
116
129
  if registry is None:
117
130
  registry = default_registry
@@ -128,6 +141,7 @@ def torch_op(
128
141
  else:
129
142
  assert isinstance(func, FunctionType)
130
143
  processed_func = onnxscript.script(opset=custom_opset)(func)
144
+ processed_func.experimental_traceable = traceable
131
145
 
132
146
  assert registry is not None
133
147
  for name_ in _check_and_normalize_names(name):
@@ -69,7 +69,7 @@ def sample_inputs_conv3d(op_info, device, dtype, requires_grad, **kwargs):
69
69
  (32,),
70
70
  {
71
71
  "stride": (3, 3, 3),
72
- "padding": 2,
72
+ "padding": (2, 2, 2),
73
73
  "dilation": (1, 1, 1),
74
74
  "groups": 1,
75
75
  },
@@ -190,21 +190,20 @@ def sample_inputs_convolution(op_info, device, dtype, requires_grad, **kwargs):
190
190
  )
191
191
 
192
192
 
193
- def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
194
- del self # Unused
193
+ def _prepare_data_for_fft_ops(device, dtype, requires_grad=False):
195
194
  # Adapted from https://github.com/pytorch/pytorch/blob/01069ad4be449f376cf88a56d842b8eb50f6e9b6/torch/testing/_internal/opinfo/core.py#L2448C1-L2541C79
196
195
  is_fp16_or_chalf = dtype in (torch.complex32, torch.half)
197
196
  if not is_fp16_or_chalf:
198
- nd_tensor = functools.partial(
197
+ oned_tensor = functools.partial(
199
198
  opinfo_core.make_tensor,
200
- (S, S + 1, S + 2),
199
+ (31,),
201
200
  device=device,
202
201
  dtype=dtype,
203
202
  requires_grad=requires_grad,
204
203
  )
205
- oned_tensor = functools.partial(
204
+ nd_tensor = functools.partial(
206
205
  opinfo_core.make_tensor,
207
- (31,),
206
+ (S, S + 1, S + 2),
208
207
  device=device,
209
208
  dtype=dtype,
210
209
  requires_grad=requires_grad,
@@ -214,18 +213,18 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
214
213
  high = None
215
214
  shapes = ((2, 8, 9), (33,))
216
215
 
217
- nd_tensor = functools.partial(
216
+ oned_tensor = functools.partial(
218
217
  opinfo_core.make_tensor,
219
- shapes[0],
218
+ shapes[1],
220
219
  device=device,
221
220
  low=low,
222
221
  high=high,
223
222
  dtype=dtype,
224
223
  requires_grad=requires_grad,
225
224
  )
226
- oned_tensor = functools.partial(
225
+ nd_tensor = functools.partial(
227
226
  opinfo_core.make_tensor,
228
- shapes[1],
227
+ shapes[0],
229
228
  device=device,
230
229
  low=low,
231
230
  high=high,
@@ -233,6 +232,13 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
233
232
  requires_grad=requires_grad,
234
233
  )
235
234
 
235
+ return oned_tensor, nd_tensor
236
+
237
+
238
+ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
239
+ del self # Unused
240
+ oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
241
+
236
242
  for normalization, forward in itertools.product((0, 1, 2), (True, False)):
237
243
  # 1-D
238
244
  yield opinfo_core.SampleInput(
@@ -252,6 +258,29 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
252
258
  )
253
259
 
254
260
 
261
+ def sample_inputs__fft_r2c(self, device, dtype, requires_grad=False, **_):
262
+ del self # Unused
263
+ oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
264
+
265
+ for normalization, one_sided in itertools.product((0, 1, 2), (True, True)):
266
+ # 1-D
267
+ yield opinfo_core.SampleInput(
268
+ oned_tensor(), dim=(0,), normalization=normalization, onesided=one_sided
269
+ )
270
+ # N-D
271
+ for dim in [
272
+ (0,),
273
+ (1,),
274
+ (2,),
275
+ (1, 2),
276
+ (0, 1),
277
+ (0, 1, 2),
278
+ ]:
279
+ yield opinfo_core.SampleInput(
280
+ nd_tensor(), dim=dim, normalization=normalization, onesided=one_sided
281
+ )
282
+
283
+
255
284
  def sample_inputs_layer_norm(op_info, device, dtype, requires_grad, **kwargs):
256
285
  del op_info # unused
257
286
  del kwargs
@@ -1336,6 +1365,25 @@ def sample_inputs__native_batch_norm_legit_no_stats(
1336
1365
  )
1337
1366
 
1338
1367
 
1368
+ def sample_inputs_reflection_pad1d(op_info, device, dtype, requires_grad, **kwargs):
1369
+ del op_info
1370
+ del kwargs
1371
+
1372
+ cases: tuple = ( # ignore
1373
+ ((2, 3), (1, 2)),
1374
+ ((4, 5), (0, 1)),
1375
+ ((6, 7), (1, 1)),
1376
+ ((8, 9), (1, 0)),
1377
+ )
1378
+
1379
+ make_inp = opinfo_core.partial(
1380
+ torch.testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
1381
+ )
1382
+
1383
+ for shape, pad in cases:
1384
+ yield opinfo_core.SampleInput(make_inp(shape), args=(pad,))
1385
+
1386
+
1339
1387
  # NOTE: How to create an OpInfo:
1340
1388
  # 1. Create a function that generates sample inputs for the op.
1341
1389
  # This function should yield SampleInputs.
@@ -1358,6 +1406,13 @@ OP_DB: List[opinfo_core.OpInfo] = [
1358
1406
  sample_inputs_func=sample_inputs__fft_c2c,
1359
1407
  supports_out=False,
1360
1408
  ),
1409
+ opinfo_core.OpInfo(
1410
+ "ops.aten._fft_r2c",
1411
+ aten_name="_fft_r2c",
1412
+ dtypes=common_dtype.floating_types(),
1413
+ sample_inputs_func=sample_inputs__fft_r2c,
1414
+ supports_out=False,
1415
+ ),
1361
1416
  opinfo_core.OpInfo(
1362
1417
  "ops.aten._local_scalar_dense",
1363
1418
  aten_name="_local_scalar_dense",
@@ -1394,7 +1449,7 @@ OP_DB: List[opinfo_core.OpInfo] = [
1394
1449
  supports_out=False,
1395
1450
  ),
1396
1451
  opinfo_core.OpInfo(
1397
- "nn.functional.conv3d",
1452
+ "ops.aten.conv3d",
1398
1453
  aten_name="conv3d",
1399
1454
  dtypes=common_dtype.floating_and_complex_types_and(torch.int64, torch.bfloat16),
1400
1455
  sample_inputs_func=sample_inputs_conv3d,
@@ -1407,6 +1462,13 @@ OP_DB: List[opinfo_core.OpInfo] = [
1407
1462
  sample_inputs_func=sample_inputs_convolution,
1408
1463
  supports_out=False,
1409
1464
  ),
1465
+ opinfo_core.OpInfo(
1466
+ "ops.aten.reflection_pad1d",
1467
+ aten_name="ops.aten.reflection_pad1d",
1468
+ dtypes=common_dtype.floating_and_complex_types_and(torch.int64, torch.bfloat16),
1469
+ sample_inputs_func=sample_inputs_reflection_pad1d,
1470
+ supports_out=False,
1471
+ ),
1410
1472
  opinfo_core.OpInfo(
1411
1473
  "ops.aten.index.Tensor",
1412
1474
  aten_name="index.Tensor",