Repository: pyro-ppl/pyro

Issue:
`poutine.block` does not work in Python 3.10
### Reproducible code

```python
from pyro import poutine

class A:
    @poutine.block
    def run(self):
        return 1

a = A()
a.run()  # error
```

This causes the issue https://github.com/pyro-ppl/pyro/issues/3018
[Support for Python 3.10] MCMC example in documentation does not work: `AttributeError: __enter__`
I am following the MCMC example in the documentation: <https://docs.pyro.ai/en/stable/mcmc.html#nuts>

```python
import pyro
import pyro.distributions as dist
from pyro.infer import MCMC, NUTS
import torch

true_coefs = torch.tensor([1., 2., 3.])
data = torch.randn(2000, 3)
dim = 3
labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()

def model(data):
    coefs_mean = torch.zeros(dim)
    coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(3)))
    y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
    return y

nuts_kernel = NUTS(model, adapt_step_size=True)
mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=300)
mcmc.run(data)
```

But it raises an error:

```
Traceback (most recent call last):
  File "/home/ayaka/Projects/test/main.py", line 19, in <module>
    mcmc.run(data)
  File "/home/ayaka/venv/lib/python3.10/site-packages/pyro/poutine/messenger.py", line 11, in _context_wrap
    with context:
AttributeError: __enter__
```

Pyro version: 1.8.0


Discussion:

Hi @ayaka14732, I just tested the code in colab and it is working fine. Maybe this issue is specific to Python 3.10?
Thanks! I tested it on Python 3.8 and it works now.
Hi @fehiepsi, Looks like this issue is also affecting some of our Ax users: https://github.com/facebook/Ax/issues/893. Are there any plans to look into this issue and have Pyro support Python 3.10? Thanks!
@dme65 Could you ping me when PyTorch [supports python 3.10](https://github.com/pytorch/pytorch/issues/66424)? I will try to fix the issue then. Thanks!
For me on windows pytorch now installs through conda in python 3.10 @fehiepsi!

I could not for the life of me figure out why after installing pytorch using the default conda install and then pyro from pip I was getting this error. I guess pyro does not check the python version when installing through pip.
> pyro does not check the python version when installing through pip

That's right. This might be an Pyro issue, not an PyTorch issue. Looking at the code, it is not clear to me why the error happens. Hopefully the issue will be resolved after PyTorch officially supports Python 3.10. (I think they set the milestone PyTorch 1.12.0, which is the next release, for it.)
@nabriis that sounds rough! 😬 

This is a PyTorch issue. I believe all of Pyro is compatible with Python 3.10, and is forward compatible with PyTorch 1.12 which will be compatible with Python 3.10.  It looks like PyTorch [only checks the min python version](https://github.com/pytorch/pytorch/blob/master/setup.py#L212), and doesn't check the max version.  I believe the right solution here is to file a PyTorch issue requesting they check both min and max Python versions.
> @nabriis that sounds rough! 😬
> 
> This is a PyTorch issue. I believe all of Pyro is compatible with Python 3.10, and is forward compatible with PyTorch 1.12 which will be compatible with Python 3.10. It looks like PyTorch [only checks the min python version](https://github.com/pytorch/pytorch/blob/master/setup.py#L212), and doesn't check the max version. I believe the right solution here is to file a PyTorch issue requesting they check both min and max Python versions.

I wonder how to file this as a pytorch issue though?

Because for what I used pytorch for it worked fine on 3.10. Only when using pyro I got the error __enter__ error. No idea what causes it or even where in the code.

Resolving patch:
diff --git a/pyro/poutine/messenger.py b/pyro/poutine/messenger.py
--- a/pyro/poutine/messenger.py
+++ b/pyro/poutine/messenger.py
@@ -18,6 +18,14 @@ class _bound_partial(partial):
     support class methods as arguments to handlers.
     """
 
+    # Use '__slots__' for func to avoid the issue
+    #   `_bound_partial(_bound_partial(f)).func is f`
+    # in Python 3.10.
+    __slots__ = "func"
+
+    def __init__(self, func):
+        self.func = func
+
     def __get__(self, instance, owner):
         if instance is None:
             return self


Test patch:
diff --git a/tests/poutine/test_poutines.py b/tests/poutine/test_poutines.py
--- a/tests/poutine/test_poutines.py
+++ b/tests/poutine/test_poutines.py
@@ -968,3 +968,13 @@ def model():
 
     with poutine.mask(mask=False):
         model()
+
+
+def test_block_class_method():
+    class A:
+        @poutine.block
+        def run(self):
+            return 1
+
+    a = A()
+    a.run()