message: tidy up doctest
diff: @@ -139,23 +139,21 @@ def custom_gradient(f=None):
  the same number of variables. We take the function `z = x * y` as an example.
  >>> @tf.custom_gradient
  - def bar(x, y):
  - def grad(upstream):
  - dz_dx = y
  - dz_dy = x
  - return upstream * dz_dx, upstream * dz_dy
  -
  - z = x * y
  -
  - return z, grad
  -
  - >>> x = tf.constant(2.0, dtype=tf.float32)
  - >>> y = tf.constant(3.0, dtype=tf.float32)
  -
  - >>> with tf.GradientTape(persistent=True) as tape:
  - tape.watch(x)
  - tape.watch(y)
  - z = bar(x, y)
  + ... def bar(x, y):
  + ... def grad(upstream):
  + ... dz_dx = y
  + ... dz_dy = x
  + ... return upstream * dz_dx, upstream * dz_dy
  + ... z = x * y
  + ... return z, grad
  + ...
  + ... x = tf.constant(2.0, dtype=tf.float32)
  + ... y = tf.constant(3.0, dtype=tf.float32)
  + ...
  + ... with tf.GradientTape(persistent=True) as tape:
  + ... tape.watch(x)
  + ... tape.watch(y)
  + ... z = bar(x, y)
  >>> z
  6
  