Add opt-einsum to CI (pytorch#85574)

Depends on pytorch#84890.

This PR adds opt_einsum to CI, enabling path optimization for the multi-input case. It also updates the installation sites to install torch with einsum, but those are mostly to make sure it would work on the user's end (as opt-einsum would have already been installed in the docker or in prior set up steps).

This PR also updates the windows build_pytorch.bat script to use the same bdist_wheel and install commands as on Linux, replacing the `setup.py install` that'll become deprecated.
Pull Request resolved: pytorch#85574
Approved by: https://github.com/huydhn, https://github.com/soulitzer

.circleci/docker/requirements-ci.txt

@@ -124,6 +124,11 @@ numba==0.55.2 ; python_version == "3.10"
 #Pinned versions: 1.9.0
 #test that import:
 
+opt-einsum==3.3
+#Description: Python library to optimize tensor contraction order, used in einsum
+#Pinned versions: 3.3
+#test that import: test_linalg.py
+
 #pillow
 #Description:  Python Imaging Library fork
 #Pinned versions:

.github/workflows/_docs.yml

@@ -118,7 +118,7 @@ jobs:
             -w /var/lib/jenkins/workspace \
             "${DOCKER_IMAGE}"
           )
-          docker exec -t "${container_name}" bash -c "sudo chown -R jenkins . && pip install dist/*.whl && ./.circleci/scripts/${DOCS_TYPE}_doc_push_script.sh"
+          docker exec -t "${container_name}" bash -c "sudo chown -R jenkins . && pip install $(echo dist/*.whl)[opt-einsum] && ./.circleci/scripts/${DOCS_TYPE}_doc_push_script.sh"
 
       - name: Chown workspace
         uses: ./.github/actions/chown-workspace

.github/workflows/_linux-test.yml

@@ -189,7 +189,7 @@ jobs:
             -w /var/lib/jenkins/workspace \
             "${DOCKER_IMAGE}"
           )
-          docker exec -t "${container_name}" sh -c "pip install dist/*.whl && ${TEST_COMMAND}"
+          docker exec -t "${container_name}" sh -c "pip install $(echo dist/*.whl)[opt-einsum] && ${TEST_COMMAND}"
 
       - name: Get workflow job id
         id: get-job-id

.github/workflows/_mac-test.yml

@@ -142,7 +142,7 @@ jobs:
           export PR_BODY="${PR_BODY//[\'\"]}"
           arch
 
-          ${CONDA_RUN} python3 -mpip install dist/*.whl
+          ${CONDA_RUN} python3 -mpip install $(echo dist/*.whl)[opt-einsum]
           ${CONDA_RUN} .jenkins/pytorch/macos-test.sh
 
       - name: Get workflow job id

.jenkins/pytorch/build-asan.sh

@@ -26,7 +26,7 @@ CC="clang" CXX="clang++" LDSHARED="clang --shared" \
   CFLAGS="-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all -fsanitize-address-use-after-scope -shared-libasan" \
   USE_ASAN=1 USE_CUDA=0 USE_MKLDNN=0 \
   python setup.py bdist_wheel
-  python -mpip install dist/*.whl
+  python -mpip install "$(echo dist/*.whl)[opt-einsum]"
 
 # Test building via the sdist source tarball
 python setup.py sdist

.jenkins/pytorch/build.sh

@@ -232,7 +232,7 @@ else
     else
       python setup.py bdist_wheel
     fi
-    python -mpip install dist/*.whl
+    python -mpip install "$(echo dist/*.whl)[opt-einsum]"
 
     # TODO: I'm not sure why, but somehow we lose verbose commands
     set -x

.jenkins/pytorch/macos-test.sh

@@ -19,7 +19,8 @@ pip install "unittest-xml-reporting<=3.2.0,>=2.0.0" \
   pytest-shard \
   pytest-rerunfailures \
   "xdoctest==1.0.2" \
-  "pygments==2.12.0"
+  "pygments==2.12.0" \
+  "opt-einsum>=3.3"
 
 if [ -z "${CI}" ]; then
   rm -rf "${WORKSPACE_DIR}"/miniconda3/lib/python3.6/site-packages/torch*

.jenkins/pytorch/win-test-helpers/build_pytorch.bat

@@ -140,7 +140,7 @@ if not x%BUILD_ENVIRONMENT:cuda11=%==x%BUILD_ENVIRONMENT% (
    set BUILD_SPLIT_CUDA=ON
 )
 
-python setup.py install --cmake && sccache --show-stats && (
+python setup.py bdist_wheel && sccache --show-stats && python -c "import os, glob; os.system('python -mpip install ' + glob.glob('dist/*.whl')[0] + '[opt-einsum]')" (
   if "%BUILD_ENVIRONMENT%"=="" (
     echo NOTE: To run `import torch`, please make sure to activate the conda environment by running `call %CONDA_PARENT_DIR%\Miniconda3\Scripts\activate.bat %CONDA_PARENT_DIR%\Miniconda3` in Command Prompt before running Git Bash.
   ) else (

.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat

@@ -36,7 +36,7 @@ popd
 =======
 :: Pin unittest-xml-reporting to freeze printing test summary logic, related: https://github.com/pytorch/pytorch/issues/69014
 
-pip install "ninja==1.10.0.post1" future "hypothesis==5.35.1" "expecttest==0.1.3" "librosa>=0.6.2" "scipy==1.6.3" psutil pillow "unittest-xml-reporting<=3.2.0,>=2.0.0" pytest pytest-xdist pytest-shard pytest-rerunfailures "xdoctest==1.0.2" "pygments==2.12.0"
+pip install "ninja==1.10.0.post1" future "hypothesis==5.35.1" "expecttest==0.1.3" "librosa>=0.6.2" "scipy==1.6.3" psutil pillow "unittest-xml-reporting<=3.2.0,>=2.0.0" pytest pytest-xdist pytest-shard pytest-rerunfailures "xdoctest==1.0.2" "pygments==2.12.0" "opt-einsum>=3.3"
 if errorlevel 1 exit /b
 if not errorlevel 0 exit /b
 

setup.py

@@ -970,6 +970,10 @@ def main():
         'typing_extensions',
     ]
 
+    extras_require = {
+        'opt-einsum': ['opt-einsum>=3.3']
+    }
+
     # Parse the command line and check the arguments before we proceed with
     # building deps and setup. We need to set values so `--help` works.
     dist = Distribution()
@@ -989,10 +993,6 @@ def main():
 
     install_requires += extra_install_requires
 
-    extras_require = {
-        'opt-einsum': ['opt-einsum>=3.3']
-    }
-
     # Read in README.md for our long_description
     with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f:
         long_description = f.read()

torch/functional.py

@@ -16,10 +16,11 @@
 from torch import _VF
 
 # Set a global declaring that we have opt_einsum
-from importlib.util import find_spec as _find_spec
-if _find_spec('opt_einsum') is not None:
+# Reasons for not being able to import include opt_einsum not being found
+# OR numpy (an opt_einsum dependency) not being found
+try:
     import opt_einsum as _opt_einsum  # type: ignore[import]
-else:
+except ModuleNotFoundError:
     _opt_einsum = None