Repository: pymc-devs/pymc

Issue:
Cannot unpickle certain step samplers when `mp_ctx=spawn` or `forkserver`
### Description

This error came up in https://github.com/pymc-devs/pymc-extras/pull/542. It is caused by `PointFunc` recursively trying to access `self.f` during unpickling. MWE:

```py
import pymc as pm
with pm.Model() as m:
    x = pm.Categorical('x', logit_p=[1., 1., 1., 1.]))
    idata = pm.sample(step=pm.CategoricalGibbsMetropolis([x]), mp_ctx='spawn')
```

A simple fix would add a check that f has been set inside `PointFunc.__getattr__` 

Discussion hints:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 92.98%. Comparing base [(`ae43026`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/ae43026767209d4a78ce812e08b20210dacb8442?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) to head [(`55db526`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/55db526134e56680366c39e066f353421cefd950?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs).
> Report is 6 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@            Coverage Diff             @@
##             main    #7858      +/-   ##
==========================================
+ Coverage   92.92%   92.98%   +0.06%     
==========================================
  Files         107      108       +1     
  Lines       18299    18328      +29     
==========================================
+ Hits        17004    17043      +39     
+ Misses       1295     1285      -10     
```

| [Files with missing lines](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/pytensorf.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858?src=pr&el=tree&filepath=pymc%2Fpytensorf.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9weXRlbnNvcmYucHk=) | `90.02% <100.00%> (ø)` | |

... and [20 files with indirect coverage changes](https://app.codecov.io/gh/pymc-devs/pymc/pull/7858/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

</details>
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Solution patch:
diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py
index f1d69c928..572241e49 100644
--- a/pymc/pytensorf.py
+++ b/pymc/pytensorf.py
@@ -567,10 +567,8 @@ class PointFunc:
     def __call__(self, state):
         return self.f(**state)
 
-    def __getattr__(self, item):
-        """Allow access to the original function attributes."""
-        # This is only reached if `__getattribute__` fails.
-        return getattr(self.f, item)
+    def dprint(self, **kwrags):
+        return self.f.dprint(**kwrags)
 
 
 class CallableTensor:


Test patch:
diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py
index f7efd5f6d..de7f7f1cb 100644
--- a/tests/test_pytensorf.py
+++ b/tests/test_pytensorf.py
@@ -746,7 +746,7 @@ def test_hessian_sign_change_warning(func):
     assert equal_computations([res_neg], [-res])
 
 
-def test_point_func():
+def test_point_func(capsys):
     x, y = pt.vectors("x", "y")
     outs = x * 2 + y**2
     f = compile([x, y], outs)
@@ -758,3 +758,30 @@ def test_point_func():
     dprint_res = point_f.dprint(file="str")
     expected_dprint_res = point_f.f.dprint(file="str")
     assert dprint_res == expected_dprint_res
+
+    point_f.dprint(print_shape=True)
+    captured = capsys.readouterr()
+
+    # The shape=(?,) arises because the inputs are dvector. This checks that the dprint works, and the print_shape
+    # kwargs was correctly forwarded
+    assert "shape=(?,)" in captured.out
+
+
+def test_pickle_point_func():
+    """
+    Regression test for https://github.com/pymc-devs/pymc/issues/7857
+    """
+    import cloudpickle
+
+    x, y = pt.vectors("x", "y")
+    outs = x * 2 + y**2
+    f = compile([x, y], outs)
+
+    point_f = PointFunc(f)
+    point_f_pickled = cloudpickle.dumps(point_f)
+    point_f_unpickled = cloudpickle.loads(point_f_pickled)
+
+    # Check that the function survived the round-trip
+    np.testing.assert_allclose(
+        point_f_unpickled({"y": [3], "x": [2]}), point_f({"y": [3], "x": [2]})
+    )