From 84868382ba82f316d8bd4bc6ee67419152745675 Mon Sep 17 00:00:00 2001
From: Yichao Zhou <broken.zhoug@gmail.com>
Date: Fri, 25 Dec 2015 19:59:00 -0800
Subject: [PATCH] update the docs

---
 README.md      | 8 ++++++--
 docs/index.rst | 8 ++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 1521530..06791bc 100644
--- a/README.md
+++ b/README.md
@@ -27,14 +27,18 @@ something like:
 import numpy as np
 def logprob(x, ivar):
     logp = -0.5 * np.sum(ivar * x**2)
+    return logp
+
+def logprob_grad(x, ivar):
     grad = -ivar * x
-    return logp, grad
+    return grad
 ```
 
 ```
 from pyhmc import hmc
 ivar = 1. / np.random.rand(5)
-samples = hmc(logprob, x0=np.random.randn(5), args=(ivar,), n_samples=1e4)
+samples = hmc(logprob, logprob_grad, x0=np.random.randn(5),
+              args=(ivar,), n_samples=1e4)
 ```
 
 ```
diff --git a/docs/index.rst b/docs/index.rst
index 1bb8f03..db4377f 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -23,8 +23,11 @@ If you wanted to draw samples from a 5 dimensional Gaussian, you would do someth
     import numpy as np
     def logprob(x, ivar):
         logp = -0.5 * np.sum(ivar * x**2)
+        return logp
+
+    def logprob_grad(x, ivar):
         grad = -ivar * x
-        return logp, grad
+        return grad
 
 
 .. code-block:: python
@@ -32,7 +35,8 @@ If you wanted to draw samples from a 5 dimensional Gaussian, you would do someth
     # run the sampler
     from pyhmc import hmc
     ivar = 1. / np.random.rand(5)
-    samples = hmc(logprob, x0=np.random.randn(5), args=(ivar,), n_samples=1e4)
+    samples = hmc(logprob, logprob_grad, x0=np.random.randn(5),
+                  args=(ivar,), n_samples=1e4)
 
 .. code-block:: python