Add compute_grad option to model.logp_dlogp_function

Author: aseyboldt

pymc3/model.py

@@ -597,6 +597,8 @@ class ValueGradFunction:
         See `numpy.can_cast` for a description of the options.
         Keep in mind that we cast the variables to the array *and*
         back from the array dtype to the variable dtype.
+    compute_grads: bool, default=True
+        If False, return only the logp, not the gradient.
     kwargs
         Extra arguments are passed on to `theano.function`.
 
@@ -611,7 +613,15 @@ class ValueGradFunction:
     """
 
     def __init__(
-        self, costs, grad_vars, extra_vars=None, dtype=None, casting="no", **kwargs
+        self,
+        costs,
+        grad_vars,
+        extra_vars=None,
+        *,
+        dtype=None,
+        casting="no",
+        compute_grads=True,
+        **kwargs
     ):
         from .distributions import TensorType
 
@@ -651,13 +661,13 @@ def __init__(
         for var in self._grad_vars:
             if not np.can_cast(var.dtype, self.dtype, casting):
                 raise TypeError(
-                    "Invalid dtype for variable %s. Can not "
-                    "cast to %s with casting rule %s." % (var.name, self.dtype, casting)
+                    f"Invalid dtype for variable {var.name}. Can not "
+                    f"cast to {self.dtype} with casting rule {casting}."
                 )
             if not np.issubdtype(var.dtype, np.floating):
                 raise TypeError(
-                    "Invalid dtype for variable %s. Must be "
-                    "floating point but is %s." % (var.name, var.dtype)
+                    f"Invalid dtype for variable {var.name}. Must be "
+                    f"floating point but is {var.dtype}."
                 )
 
         givens = []
@@ -677,13 +687,17 @@ def __init__(
             self._cost, grad_vars, self._ordering.vmap
         )
 
-        grad = tt.grad(self._cost_joined, self._vars_joined)
-        grad.name = "__grad"
+        if compute_grads:
+            grad = tt.grad(self._cost_joined, self._vars_joined)
+            grad.name = "__grad"
+            outputs = [self._cost_joined, grad]
+        else:
+            outputs = self._cost_joined
 
         inputs = [self._vars_joined]
 
         self._theano_function = theano.function(
-            inputs, [self._cost_joined, grad], givens=givens, **kwargs
+            inputs, outputs, givens=givens, **kwargs
         )
 
     def set_weights(self, values):
@@ -723,12 +737,12 @@ def __call__(self, array, grad_out=None, extra_vars=None):
         else:
             out = grad_out
 
-        logp, dlogp = self._theano_function(array)
+        output = self._theano_function(array)
         if grad_out is None:
-            return logp, dlogp
+            return output
         else:
-            np.copyto(out, dlogp)
-            return logp
+            np.copyto(out, output[1])
+            return output[0]
 
     @property
     def profile(self):

pymc3/tests/test_model.py

@@ -393,8 +393,18 @@ def test_tempered_logp_dlogp():
     func_temp = model.logp_dlogp_function(tempered=True)
     func_temp.set_extra_values({})
 
+    func_nograd = model.logp_dlogp_function(compute_grads=False)
+    func_nograd.set_extra_values({})
+
+    func_temp_nograd = model.logp_dlogp_function(
+        tempered=True, compute_grads=False
+    )
+    func_temp_nograd.set_extra_values({})
+
     x = np.ones(func.size, dtype=func.dtype)
     assert func(x) == func_temp(x)
+    assert func_nograd(x) == func(x)[0]
+    assert func_temp_nograd(x) == func(x)[0]
 
     func_temp.set_weights(np.array([0.], dtype=func.dtype))
     func_temp_nograd.set_weights(np.array([0.], dtype=func.dtype))
@@ -408,3 +418,6 @@ def test_tempered_logp_dlogp():
     func_temp_nograd.set_weights(np.array([0.5], dtype=func.dtype))
     npt.assert_allclose(func(x)[0], 4 / 3 * func_temp(x)[0])
     npt.assert_allclose(func(x)[1], func_temp(x)[1])
+
+    npt.assert_allclose(func_nograd(x), func(x)[0])
+    npt.assert_allclose(func_temp_nograd(x), func_temp(x)[0])