Fix formatting of lax.while_loop and lax.fori_loop doc comments.

google · Feb 18, 2019 · 97fb6f1 · 97fb6f1
1 parent a44d248
commit 97fb6f1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/lax.py b/jax/lax.py
@@ -427,7 +427,7 @@ def while_loop(cond_fun, body_fun, init_val):
   Returns:
     The output from the final iteration of body_fun, of type `T`.
 
-  The semantics of `while_loop` are given by this Python implementation:
+  The semantics of `while_loop` are given by this Python implementation::
 
     def while_loop(cond_fun, body_fun, init_val):
       val = init_val
@@ -648,7 +648,7 @@ def fori_loop(lower, upper, body_fun, init_val):
   Returns:
     Loop value from the final iteration, of type T.
 
-  The semantics of `fori_loop` are given by this Python implementation:
+  The semantics of `fori_loop` are given by this Python implementation::
 
     def fori_loop(lower, upper, body_fun, init_val):
       val = init_val