Add docstrings to various logp related methods

Author: michaelosthege

pymc/model.py

@@ -651,23 +651,53 @@ def compile_logp(
         jacobian: bool = True,
         sum: bool = True,
     ):
-        """Compiled log probability density function"""
+        """Compiled log probability density function.
+
+        Parameters
+        ----------
+        vars: list of random variables or potential terms, optional
+            Compute the gradient with respect to those variables. If None, use all
+            free and observed random variables, as well as potential terms in model.
+        jacobian:
+            Whether to include jacobian terms in logprob graph. Defaults to True.
+        sum:
+            Whether to sum all logp terms or return elemwise logp for each variable.
+            Defaults to True.
+        """
         return self.model.compile_fn(self.logpt(vars=vars, jacobian=jacobian, sum=sum))
 
     def compile_dlogp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
     ):
-        """Compiled log probability density gradient function"""
+        """Compiled log probability density gradient function.
+
+        Parameters
+        ----------
+        vars: list of random variables or potential terms, optional
+            Compute the gradient with respect to those variables. If None, use all
+            free and observed random variables, as well as potential terms in model.
+        jacobian:
+            Whether to include jacobian terms in logprob graph. Defaults to True.
+        """
         return self.model.compile_fn(self.dlogpt(vars=vars, jacobian=jacobian))
 
     def compile_d2logp(
         self,
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
     ):
-        """Compiled log probability density hessian function"""
+        """Compiled log probability density hessian function.
+
+        Parameters
+        ----------
+        vars: list of random variables or potential terms, optional
+            Compute the gradient with respect to those variables. If None, use all
+            free and observed random variables, as well as potential terms in model.
+        jacobian:
+            Whether to include jacobian terms in logprob graph. Defaults to True.
+        """
         return self.model.compile_fn(self.d2logpt(vars=vars, jacobian=jacobian))
 
     def logpt(
@@ -747,6 +777,20 @@ def dlogpt(
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
     ) -> Variable:
+        """Gradient of the models log-probability w.r.t. ``vars``.
+
+        Parameters
+        ----------
+        vars: list of random variables or potential terms, optional
+            Compute the gradient with respect to those variables. If None, use all
+            free and observed random variables, as well as potential terms in model.
+        jacobian:
+            Whether to include jacobian terms in logprob graph. Defaults to True.
+
+        Returns
+        -------
+        dlogp graph
+        """
         if vars is None:
             value_vars = None
         else:
@@ -771,6 +815,20 @@ def d2logpt(
         vars: Optional[Union[Variable, Sequence[Variable]]] = None,
         jacobian: bool = True,
     ) -> Variable:
+        """Hessian of the models log-probability w.r.t. ``vars``.
+
+        Parameters
+        ----------
+        vars: list of random variables or potential terms, optional
+            Compute the gradient with respect to those variables. If None, use all
+            free and observed random variables, as well as potential terms in model.
+        jacobian:
+            Whether to include jacobian terms in logprob graph. Defaults to True.
+
+        Returns
+        -------
+        d²logp graph
+        """
         if vars is None:
             value_vars = None
         else:

pymc/tests/test_shared.py

@@ -27,7 +27,6 @@ def test_deterministic(self):
             data_values = np.array([0.5, 0.4, 5, 2])
             X = aesara.shared(np.asarray(data_values, dtype=aesara.config.floatX), borrow=True)
             pm.Normal("y", 0, 1, observed=X)
-            # TODO: This should assert something
             assert np.all(
                 np.isclose(model.compile_logp(sum=False)({}), st.norm().logpdf(data_values))
             )