Revert to support torch>=1.11.0 (#3242)

.github/workflows/ci.yml

@@ -54,7 +54,7 @@ jobs:
           python -m pip install --upgrade pip wheel 'setuptools!=58.5.*,<60'
           # Keep track of pyro-api master branch
           pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip
-          pip install torch==1.11.0+cpu torchvision==0.12.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+          pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
           pip install .[test]
           pip install -r docs/requirements.txt
           pip freeze
@@ -82,7 +82,7 @@ jobs:
           python -m pip install --upgrade pip wheel 'setuptools!=58.5.*,<60'
           # Keep track of pyro-api master branch
           pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip
-          pip install torch==1.11.0+cpu torchvision==0.12.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+          pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
           pip install .[test]
           pip install -r docs/requirements.txt
           # requirements for tutorials (from .[dev])

pyro/ops/provenance.py

@@ -93,6 +93,11 @@ def _track_provenance_set(x, provenance: frozenset):
 @track_provenance.register(tuple)
 @track_provenance.register(dict)
 def _track_provenance_pytree(x, provenance: frozenset):
+    # avoid max-recursion depth error for torch<=2.0
+    flat_args, _ = tree_flatten(x)
+    if not flat_args or flat_args[0] is x:
+        return x
+
     return tree_map(partial(track_provenance, provenance=provenance), x)
 
 
@@ -138,6 +143,11 @@ def _extract_provenance_set(x):
 @extract_provenance.register(tuple)
 @extract_provenance.register(dict)
 def _extract_provenance_pytree(x):
+    # avoid max-recursion depth error for torch<=2.0
+    flat_args, _ = tree_flatten(x)
+    if not flat_args or flat_args[0] is x:
+        return x, frozenset()
+
     flat_args, spec = tree_flatten(x)
     xs = []
     provenance = frozenset()

setup.py

@@ -102,7 +102,7 @@
         "numpy>=1.7",
         "opt_einsum>=2.3.2",
         "pyro-api>=0.1.1",
-        "torch>=2.0.1",
+        "torch>=1.11.0",
         "tqdm>=4.36",
     ],
     extras_require={
@@ -112,6 +112,7 @@
             "black>=21.4b0",
             "nbval",
             "pytest-cov",
+            "pytest-xdist",
             "pytest>=5.0",
             "ruff",
             "scipy>=1.1",