Update on "[ONNX] Add test_fx_op_consistency.py"

Should be merged after #99434

<!--
copilot:all
-->
### <samp>🤖 Generated by Copilot at f59c428</samp>

### Summary
📝🔄🚀

<!--
1.  📝 - This emoji represents the improvement of the documentation and type annotation of the ONNX exporter and its tests.
2.  🔄 - This emoji represents the renaming of the `dont_care` function to `skip` and the update of the references and documentation. This change reflects a refactoring and improvement of the code quality and readability.
3.  🚀 - This emoji represents the enhancement of the ONNX exporter's ability to handle more input and output types for PyTorch models. This change increases the performance and functionality of the exporter.
-->
This pull request enhances the ONNX exporter and its tests to handle more input and output types for PyTorch models. It improves the type annotation, tolerance handling, and documentation of the exporter and its test functions. It also renames a test function to make the code more consistent and clear.

> _To export PyTorch models with ONNX_
> _We need to handle various contexts_
> _We improved the annotation_
> _And the tolerance function_
> _And renamed `dont_care` to `skip` for the tests_

### Walkthrough
*  Expand the `_InputArgsType` type annotation to include int, float, and bool types, in addition to torch.Tensor, Sequence, and Mapping, to support more types of inputs for PyTorch models ([link](https://github.com/pytorch/pytorch/pull/99465/files?diff=unified&w=0#diff-c8fa56eefd7f98fb4f9739d57df57f02ede77e28528133736010a6d06651ebcbL44-R46), [link](https://github.com/pytorch/pytorch/pull/99465/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L197-R199), [link](https://github.com/pytorch/pytorch/pull/99465/files?diff=unified&w=0#diff-0795f54fd1f38cfbf2c4a863a4efc9f40f2ea020a2b1612605c361b8d8d35862L291-R295)).



[ghstack-poisoned]

.bazelrc

@@ -69,10 +69,6 @@ build --per_file_copt='^//.*\.(cpp|cc)$'@-Werror=all
 # The following warnings come from -Wall. We downgrade them from error
 # to warnings here.
 #
-# sign-compare has a tremendous amount of violations in the
-# codebase. It will be a lot of work to fix them, just disable it for
-# now.
-build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-sign-compare
 # We intentionally use #pragma unroll, which is compiler specific.
 build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-error=unknown-pragmas
 

.ci/docker/ci_commit_pins/triton.txt

@@ -1 +1 @@
-46672772b46b103db7341c9e10fbad7f643557d4
+7d1a95b04654ff9c216afe08a454ad0822f05370

.ci/docker/common/install_onnx.sh

@@ -27,7 +27,7 @@ pip_install \
   transformers==4.25.1
 
 # TODO: change this when onnx-script is on testPypi
-pip_install "onnx-script@git+https://github.com/microsoft/onnx-script@1e8d764a9be04323d7171e4d5f511332790cb809"
+pip_install "onnxscript@git+https://github.com/microsoft/onnxscript@99a665546d374af872cf136137b5d8334389b78d"
 
 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers
 # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/

.ci/docker/requirements-ci.txt

@@ -263,3 +263,8 @@ jinja2==3.1.2
 #Description: jinja2 template engine
 #Pinned versions: 3.1.2
 #test that import:
+
+pytest-cpp==2.3.0
+#Description: This is used by pytest to invoke C++ tests
+#Pinned versions: 2.3.0
+#test that import:

.ci/pytorch/macos-build.sh

@@ -47,11 +47,7 @@ compile_arm64() {
 }
 
 compile_x86_64() {
-  USE_DISTRIBUTED=0 WERROR=1 python setup.py bdist_wheel
-}
-
-compile_x86_64() {
-  USE_DISTRIBUTED=0 WERROR=1 python setup.py bdist_wheel
+  USE_DISTRIBUTED=0 WERROR=1 python setup.py bdist_wheel --plat-name=macosx_10_9_x86_64
 }
 
 build_lite_interpreter() {

.ci/pytorch/macos-common.sh

@@ -9,7 +9,7 @@ sysctl -a | grep machdep.cpu
 
 # These are required for both the build job and the test job.
 # In the latter to test cpp extensions.
-export MACOSX_DEPLOYMENT_TARGET=10.9
+export MACOSX_DEPLOYMENT_TARGET=11.0
 export CXX=clang++
 export CC=clang
 

.ci/pytorch/macos-test.sh

@@ -25,7 +25,7 @@ setup_test_python() {
   # using the address associated with the loopback interface.
   export GLOO_SOCKET_IFNAME=lo0
   echo "Ninja version: $(ninja --version)"
-  echo "Python version: $(python --version)"
+  echo "Python version: $(which python) ($(python --version))"
 
   # Increase default limit on open file handles from 256 to 1024
   ulimit -n 1024

.ci/pytorch/multigpu-test.sh

@@ -35,7 +35,8 @@ time python test/run_test.py --verbose -i distributed/_shard/sharded_tensor/ops/
 time python test/run_test.py --verbose -i distributed/_shard/sharded_tensor/ops/test_init
 time python test/run_test.py --verbose -i distributed/_shard/sharded_optim/test_sharded_optim
 
-# DTensor/Random Ops tests
+# DTensor tests
+time python test/run_test.py --verbose -i distributed/_tensor/test_device_mesh.py
 time python test/run_test.py --verbose -i distributed/_tensor/test_random_ops.py
 
 # DTensor/TP tests

.ci/pytorch/print_sccache_log.py

@@ -13,5 +13,5 @@
         'CheckSymbolExists.c',
         'test_compilation_error_formatting',
     ]
-    if all([keyword not in line for keyword in ignored_keywords]):
+    if all(keyword not in line for keyword in ignored_keywords):
         print(line)

.ci/pytorch/test.sh

@@ -251,6 +251,11 @@ test_inductor() {
   python test/run_test.py --inductor --include test_modules test_ops test_ops_gradients test_torch --verbose
   # Do not add --inductor for the following inductor unit tests, otherwise we will fail because of nested dynamo state
   python test/run_test.py --include inductor/test_torchinductor inductor/test_torchinductor_opinfo --verbose
+
+  # docker build uses bdist_wheel which does not work with test_aot_inductor
+  # TODO: need a faster way to build
+  BUILD_AOT_INDUCTOR_TEST=1 python setup.py develop
+  LD_LIBRARY_PATH="$TORCH_LIB_DIR $TORCH_BIN_DIR"/test_aot_inductor
 }
 
 # "Global" flags for inductor benchmarking controlled by TEST_CONFIG
@@ -296,33 +301,42 @@ test_perf_for_dashboard() {
     # Run accuracy test for inductor with different configs
     # --disable-cudagraphs is the default inductor behavior
     # TODO: update here once cudagraphs is turned on as default
-    if [[ "${TEST_CONFIG}" != *max_autotune* ]]; then
+    if [[ "${TEST_CONFIG}" == *max_autotune* ]]; then
+      # Only run this one config for max-autotune
+      python "benchmarks/dynamo/$suite.py" \
+          --accuracy --"$mode" --"$dtype" --backend "$backend" "$@" \
+          --output "$TEST_REPORTS_DIR/${backend}_max_autotune_${suite}_${dtype}_${mode}_cuda_accuracy.csv"
+    else
       python "benchmarks/dynamo/$suite.py" \
           --accuracy --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \
           --output "$TEST_REPORTS_DIR/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_cuda_accuracy.csv"
       python "benchmarks/dynamo/$suite.py" \
-          --accuracy --"$mode" --"$dtype" --backend "$backend" --dynamic-shapes --dynamic-batch-only --disable-cudagraphs "$@" \
+          --accuracy --"$mode" --"$dtype" --backend "$backend" "$@" \
+          --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_cuda_accuracy.csv"
+      python "benchmarks/dynamo/$suite.py" \
+          --accuracy --"$mode" --"$dtype" --backend "$backend" --dynamic-shapes \
+          --dynamic-batch-only --disable-cudagraphs "$@" \
           --output "$TEST_REPORTS_DIR/${backend}_dynamic_${suite}_${dtype}_${mode}_cuda_accuracy.csv"
     fi
-    # Only test this one config for max-autotune
-    python "benchmarks/dynamo/$suite.py" \
-        --accuracy --"$mode" --"$dtype" --backend "$backend" "$@" \
-        --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_cuda_accuracy.csv"
 
     # Run performance test
-    if [[ "${TEST_CONFIG}" != *max_autotune* ]]; then
+    if [[ "${TEST_CONFIG}" == *max_autotune* ]]; then
+      # Only run this one config for max-autotune
+      python "benchmarks/dynamo/$suite.py" \
+          --performance --cold-start-latency --"$mode" --"$dtype" --backend "$backend" "$@" \
+          --output "$TEST_REPORTS_DIR/${backend}_max_autotune_${suite}_${dtype}_${mode}_cuda_performance.csv"
+    else
       python "benchmarks/dynamo/$suite.py" \
           --performance --cold-start-latency --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \
           --output "$TEST_REPORTS_DIR/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_cuda_performance.csv"
+      python "benchmarks/dynamo/$suite.py" \
+          --performance --cold-start-latency --"$mode" --"$dtype" --backend "$backend" "$@" \
+          --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_cuda_performance.csv"
       python "benchmarks/dynamo/$suite.py" \
           --performance --cold-start-latency --"$mode" --"$dtype" --backend "$backend" --dynamic-shapes \
           --dynamic-batch-only --disable-cudagraphs "$@" \
           --output "$TEST_REPORTS_DIR/${backend}_dynamic_${suite}_${dtype}_${mode}_cuda_performance.csv"
     fi
-    # Only test this one config for max-autotune
-    python "benchmarks/dynamo/$suite.py" \
-        --performance --cold-start-latency --"$mode" --"$dtype" --backend "$backend" "$@" \
-        --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_cuda_performance.csv"
   done
 }
 
@@ -413,7 +427,7 @@ test_inductor_torchbench_smoketest_perf() {
   python benchmarks/dynamo/check_hf_bert_perf_csv.py -f "$TEST_REPORTS_DIR/inductor_training_smoketest.csv"
 
   # Check memory compression ratio for a few models
-  for test in hf_Albert timm_efficientdet timm_vision_transformer; do
+  for test in hf_Albert timm_vision_transformer; do
     python benchmarks/dynamo/torchbench.py --device cuda --performance --backend inductor --amp --training \
       --disable-cudagraphs --batch-size-file "$(realpath benchmarks/dynamo/torchbench_models_list.txt)" \
       --only $test --output "$TEST_REPORTS_DIR/inductor_training_smoketest_$test.csv"
@@ -542,6 +556,7 @@ test_libtorch() {
       # TODO: Consider to run static_runtime_test from $TORCH_BIN_DIR (may need modify build script)
       "$BUILD_BIN_DIR"/static_runtime_test --gtest_output=xml:$TEST_REPORTS_DIR/static_runtime_test.xml
     fi
+
     assert_git_not_dirty
   fi
 }
@@ -940,7 +955,7 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
   install_torchvision
   id=$((SHARD_NUMBER-1))
   if [[ "${TEST_CONFIG}" == *inductor_torchbench_smoketest_perf* ]]; then
-    checkout_install_torchbench hf_Bert hf_Albert timm_efficientdet timm_vision_transformer
+    checkout_install_torchbench hf_Bert hf_Albert timm_vision_transformer
     PYTHONPATH=$(pwd)/torchbench test_inductor_torchbench_smoketest_perf
   else
     checkout_install_torchbench

.ci/pytorch/win-test.sh

@@ -34,7 +34,8 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
   export PYTORCH_TESTING_DEVICE_ONLY_FOR="cuda"
 fi
 
-python -m pip install pytest-rerunfailures==10.3
+# TODO: Move both of them to Windows AMI
+python -m pip install pytest-rerunfailures==10.3 pytest-cpp==2.3.0
 
 run_tests() {
     # Run nvidia-smi if available

.circleci/scripts/binary_windows_build.sh

@@ -8,7 +8,7 @@ export CUDA_VERSION="${DESIRED_CUDA/cu/}"
 export USE_SCCACHE=1
 export SCCACHE_BUCKET=ossci-compiler-cache
 export SCCACHE_IGNORE_SERVER_IO_ERROR=1
-export VC_YEAR=2022
+export VC_YEAR=2019
 
 if [[ "${DESIRED_CUDA}" =~ cu1[1-2][0-9] ]]; then
     export BUILD_SPLIT_CUDA=ON

.circleci/scripts/binary_windows_test.sh

@@ -4,7 +4,7 @@ set -eux -o pipefail
 source "${BINARY_ENV_FILE:-/c/w/env}"
 
 export CUDA_VERSION="${DESIRED_CUDA/cu/}"
-export VC_YEAR=2022
+export VC_YEAR=2019
 
 pushd "$BUILDER_ROOT"
 

.circleci/scripts/cpp_doc_push_script.sh

@@ -20,7 +20,7 @@ echo "cpp_doc_push_script.sh: Invoked with $*"
 #       but since DOCS_INSTALL_PATH can be derived from DOCS_VERSION it's probably better to
 #       try and gather it first, just so we don't potentially break people who rely on this script
 # Argument 2: What version of the Python API docs we are building.
-version="${2:-${DOCS_VERSION:-master}}"
+version="${2:-${DOCS_VERSION:-main}}"
 if [ -z "$version" ]; then
 echo "error: cpp_doc_push_script.sh: version (arg2) not specified"
   exit 1
@@ -35,7 +35,7 @@ echo "error: cpp_doc_push_script.sh: install_path (arg1) not specified"
 fi
 
 is_main_doc=false
-if [ "$version" == "master" ]; then
+if [ "$version" == "main" ]; then
   is_main_doc=true
 fi
 

.circleci/scripts/functorch_doc_push_script.sh

@@ -24,7 +24,7 @@ popd
 git clone https://github.com/pytorch/functorch -b gh-pages --depth 1 functorch_ghpages
 pushd functorch_ghpages
 
-if [ $version == "master" ]; then
+if [ $version == "main" ]; then
   version=nightly
 fi
 

.circleci/scripts/python_doc_push_script.sh

@@ -23,7 +23,7 @@ set -ex
 #       but since DOCS_INSTALL_PATH can be derived from DOCS_VERSION it's probably better to
 #       try and gather it first, just so we don't potentially break people who rely on this script
 # Argument 2: What version of the docs we are building.
-version="${2:-${DOCS_VERSION:-master}}"
+version="${2:-${DOCS_VERSION:-main}}"
 if [ -z "$version" ]; then
 echo "error: python_doc_push_script.sh: version (arg2) not specified"
   exit 1
@@ -38,7 +38,7 @@ echo "error: python_doc_push_script.sh: install_path (arg1) not specified"
 fi
 
 is_main_doc=false
-if [ "$version" == "master" ]; then
+if [ "$version" == "main" ]; then
   is_main_doc=true
 fi
 

.clang-tidy

@@ -9,6 +9,7 @@ bugprone-*,
 -bugprone-lambda-function-name,
 -bugprone-reserved-identifier,
 -bugprone-swapped-arguments,
+clang-diagnostic-missing-prototypes,
 cppcoreguidelines-*,
 -cppcoreguidelines-avoid-do-while,
 -cppcoreguidelines-avoid-magic-numbers,

.flake8

@@ -1,4 +1,6 @@
 [flake8]
+# NOTE: **Mirror any changes** to this file the [tool.ruff] config in pyproject.toml
+# before we can fully move to use ruff
 enable-extensions = G
 select = B,C,E,F,G,P,SIM1,T4,W,B9
 max-line-length = 120
@@ -14,7 +16,7 @@ ignore =
     # these ignores are from flake8-bugbear; please fix!
     B007,B008,B017,B019,B020,B023,B024,B026,B027,B028,B903,B904,B905,B906,B907
     # these ignores are from flake8-comprehensions; please fix!
-    C407,C419,
+    C407,
     # these ignores are from flake8-logging-format; please fix!
     G100,G101,G200,G201,G202
     # these ignores are from flake8-simplify. please fix or ignore with commented reason

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -44,7 +44,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
+      wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
       # For security purposes, please check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/pt2-bug-report.yml

@@ -53,9 +53,9 @@ body:
       description: |
         Please run the following and paste the output below.
         ```sh
-        wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
+        curl -OL https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
         # For security purposes, please check the contents of collect_env.py before running it.
-        python collect_env.py
+        python3 collect_env.py
         ```
     validations:
       required: true

.github/ci_commit_pins/vision.txt

@@ -1 +1 @@
-b78d98bb152ffb9c0c0f5365f59f475c70b1784e
+0d75d9e5516f446c9c0ef93bd4ed9fea13992d06

.github/ci_commit_pins/xla.txt

@@ -1 +1 @@
-f235d4da06905b35d75879a0a9bc3034ab7385ac
+486b32e5b1a2e2964a7db4c7db25b59d879b8796

.github/merge_rules.yaml

@@ -420,6 +420,7 @@
   - gchanan
   - ezyang
   - dzhulgakov
+  - malfet
   mandatory_checks_name:
   - EasyCLA
   - Lint

.github/requirements/pip-requirements-macOS.txt

@@ -22,3 +22,4 @@ unittest-xml-reporting<=3.2.0,>=2.0.0
 xdoctest==1.1.0
 filelock==3.6.0
 sympy==1.11.1
+pytest-cpp==2.3.0

.github/scripts/test_trymerge.py

@@ -364,7 +364,7 @@ def test_cancelled_gets_ignored(self, *args: Any) -> None:
         lint_checks = [name for name in conclusions.keys() if "Lint" in name]
         self.assertTrue(len(lint_checks) > 0)
         self.assertTrue(
-            all([conclusions[name].status == "SUCCESS" for name in lint_checks])
+            all(conclusions[name].status == "SUCCESS" for name in lint_checks)
         )
 
     @mock.patch("trymerge.gh_get_pr_info", return_value=mock_gh_get_info())

.github/scripts/trymerge.py

@@ -82,10 +82,7 @@ def matches(self, job: Optional[Dict[str, Any]]) -> bool:
             and self.name in job.get("name", "")
             and job.get("failure_captures") is not None
             and all(
-                [
-                    capture in job.get("failure_captures", [])
-                    for capture in self.captures
-                ]
+                capture in job.get("failure_captures", []) for capture in self.captures
             )
         )
 
@@ -1575,7 +1572,7 @@ def insert(d: Dict[str, Dict[str, Any]], key: str, val: Dict[str, Any]) -> None:
             checks_with_classifications[name] = JobCheckState(
                 check.name, check.url, check.status, "BROKEN_TRUNK", check.job_id
             )
-        elif any([rule.matches(head_sha_job) for rule in flaky_rules]):
+        elif any(rule.matches(head_sha_job) for rule in flaky_rules):
             checks_with_classifications[name] = JobCheckState(
                 check.name, check.url, check.status, "FLAKY", check.job_id
             )
@@ -1710,11 +1707,11 @@ def categorize_checks(
     relevant_checknames = [
         name
         for name in check_runs.keys()
-        if not required_checks or any([x in name for x in required_checks])
+        if not required_checks or any(x in name for x in required_checks)
     ]
 
     for checkname in required_checks:
-        if all([checkname not in x for x in check_runs.keys()]):
+        if all(checkname not in x for x in check_runs.keys()):
             pending_checks.append((checkname, None, None))
 
     for checkname in relevant_checknames:

.github/workflows/_docs.yml

@@ -156,7 +156,7 @@ jobs:
           retention-days: 14
           s3-bucket: doc-previews
           if-no-files-found: error
-          path: pytorch.github.io/docs/master/
+          path: pytorch.github.io/docs/main/
           s3-prefix: pytorch/${{ github.event.pull_request.number }}
 
       - name: Upload C++ Docs Preview
@@ -190,7 +190,7 @@ jobs:
           retention-days: 14
           s3-bucket: doc-previews
           if-no-files-found: error
-          path: pytorch.github.io/docs/master/
+          path: pytorch.github.io/docs/main/
           s3-prefix: pytorch/pytorch/pytorch/${{ github.event.pull_request.number }}
 
       - name: Upload C++ Docs Preview (forward compatibility)