Relocate PointFunc definition for nicer type hints

Author: michaelosthege

pymc/aesaraf.py

@@ -677,6 +677,16 @@ def reshape_t(x, shape):
         return x[0]
 
 
+class PointFunc:
+    """Wraps so a function so it takes a dict of arguments instead of arguments."""
+
+    def __init__(self, f):
+        self.f = f
+
+    def __call__(self, state):
+        return self.f(**state)
+
+
 class CallableTensor:
     """Turns a symbolic variable with one input into a function that returns symbolic arguments
     with the one variable replaced with the input.

pymc/model.py

@@ -48,6 +48,7 @@
 from aesara.tensor.var import TensorConstant, TensorVariable
 
 from pymc.aesaraf import (
+    PointFunc,
     compile_pymc,
     convert_observed_data,
     gradient,
@@ -640,7 +641,7 @@ def compile_logp(
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
         sum: bool = True,
-    ):
+    ) -> PointFunc:
         """Compiled log probability density function.
 
         Parameters
@@ -660,7 +661,7 @@ def compile_dlogp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
-    ):
+    ) -> PointFunc:
         """Compiled log probability density gradient function.
 
         Parameters
@@ -677,7 +678,7 @@ def compile_d2logp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
-    ):
+    ) -> PointFunc:
         """Compiled log probability density hessian function.
 
         Parameters
@@ -1597,7 +1598,7 @@ def compile_fn(
         mode=None,
         point_fn: bool = True,
         **kwargs,
-    ) -> Union["PointFunc", Callable[[Sequence[np.ndarray]], Sequence[np.ndarray]]]:
+    ) -> Union[PointFunc, Callable[[Sequence[np.ndarray]], Sequence[np.ndarray]]]:
         """Compiles an Aesara function
 
         Parameters
@@ -1913,16 +1914,6 @@ def Point(*args, filter_model_vars=False, **kwargs) -> Dict[str, np.ndarray]:
     }
 
 
-class PointFunc:
-    """Wraps so a function so it takes a dict of arguments instead of arguments."""
-
-    def __init__(self, f):
-        self.f = f
-
-    def __call__(self, state):
-        return self.f(**state)
-
-
 def Deterministic(name, var, model=None, dims=None, auto=False):
     """Create a named deterministic variable