Allow accessing wrapped function attributes in PointFunc

Author: ricardoV94

pymc/pytensorf.py

@@ -600,6 +600,12 @@ def __init__(self, f):
     def __call__(self, state):
         return self.f(**state)
 
+    def __getattr__(self, item):
+        """Allow access to the original function attributes."""
+        if item == "f":
+            return self.f
+        return getattr(self.f, item)
+
 
 class CallableTensor:
     """Turns a symbolic variable with one input into a function that returns symbolic arguments with the one variable replaced with the input."""

tests/test_pytensorf.py

@@ -36,6 +36,7 @@
 from pymc.exceptions import NotConstantValueError
 from pymc.logprob.utils import ParameterValueError
 from pymc.pytensorf import (
+    PointFunc,
     collect_default_updates,
     compile,
     constant_fold,
@@ -780,3 +781,17 @@ def test_hessian_sign_change_warning(func):
         res_neg = func(f, vars=[x])
     res = func(f, vars=[x], negate_output=False)
     assert equal_computations([res_neg], [-res])
+
+
+def test_point_func():
+    x, y = pt.vectors("x", "y")
+    outs = x * 2 + y**2
+    f = compile([x, y], outs)
+
+    point_f = PointFunc(f)
+    np.testing.assert_allclose(point_f({"y": [3], "x": [2]}), [4 + 9])
+
+    # Check we can access other methods of the wrapped pytensor function
+    dprint_res = point_f.dprint(file="str")
+    expected_dprint_res = point_f.f.dprint(file="str")
+    assert dprint_res == expected_dprint_res