Pin the Python package to the exact PyTorch nightly.

This avoids issues where PyTorch version drift has made things
incompatible.

One caveat is that you will need to specify
`-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
--pre` on the command line for pip to know where to find the nightly
packages (there is no way around this) -- this is easiest to do by
simultaneously passing `-r requirements.txt` on the pip command line.

build_tools/build_python_wheels.sh

@@ -23,12 +23,13 @@ echo "---- CREATING VENV ----"
 python -m venv "$package_test_venv"
 VENV_PYTHON="$package_test_venv/bin/python"
 
-echo "---- INSTALLING torch ----"
-$VENV_PYTHON -m pip install -r "${repo_root}/requirements.txt"
+# Install the Torch-MLIR package.
+# Note that we also need to pass in the `-r requirements.txt` here to pick up
+# the right --find-links flag for the nightly PyTorch wheel registry.
+echo "---- INSTALLING torch-mlir and dependencies ----"
+$VENV_PYTHON -m pip install -f "$wheelhouse" --force-reinstall torch_mlir -r "${repo_root}/requirements.txt"
 echo "---- INSTALLING other deps for smoke test ----"
 $VENV_PYTHON -m pip install requests pillow
-echo "---- INSTALLING torch-mlir ----"
-$VENV_PYTHON -m pip install -f "$wheelhouse" --force-reinstall torch_mlir
 
 echo "---- RUNNING SMOKE TEST ----"
 $VENV_PYTHON "$repo_root/examples/torchscript_resnet18.py"

setup.py

@@ -39,6 +39,8 @@ from setuptools import setup, Extension
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
+import torch
+
 PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION") or "0.0.1"
 
 # Build phase discovery is unreliable. Just tell it what phases to run.
@@ -127,5 +129,10 @@ setup(
     ext_modules=[
         CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
     ],
+    install_requires=[
+        # To avoid issues with drift for each nightly build, we pin to the
+        # exact version we built against.
+        f"torch=={torch.__version__}",
+    ],
     zip_safe=False,
 )