Allow defining mode from compile_* functions

Author: ricardoV94

pymc/model/core.py

@@ -612,6 +612,7 @@ def compile_logp(
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
         sum: bool = True,
+        **compile_kwargs,
     ) -> PointFunc:
         """Compiled log probability density function.
 
@@ -626,12 +627,13 @@ def compile_logp(
             Whether to sum all logp terms or return elemwise logp for each variable.
             Defaults to True.
         """
-        return self.compile_fn(self.logp(vars=vars, jacobian=jacobian, sum=sum))
+        return self.compile_fn(self.logp(vars=vars, jacobian=jacobian, sum=sum), **compile_kwargs)
 
     def compile_dlogp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
+        **compile_kwargs,
     ) -> PointFunc:
         """Compiled log probability density gradient function.
 
@@ -643,12 +645,13 @@ def compile_dlogp(
         jacobian : bool
             Whether to include jacobian terms in logprob graph. Defaults to True.
         """
-        return self.compile_fn(self.dlogp(vars=vars, jacobian=jacobian))
+        return self.compile_fn(self.dlogp(vars=vars, jacobian=jacobian), **compile_kwargs)
 
     def compile_d2logp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
+        **compile_kwargs,
     ) -> PointFunc:
         """Compiled log probability density hessian function.
 
@@ -660,7 +663,7 @@ def compile_d2logp(
         jacobian : bool
             Whether to include jacobian terms in logprob graph. Defaults to True.
         """
-        return self.compile_fn(self.d2logp(vars=vars, jacobian=jacobian))
+        return self.compile_fn(self.d2logp(vars=vars, jacobian=jacobian), **compile_kwargs)
 
     def logp(
         self,