Fix the formatting of TensorArray. __doc__ indents anything after the first line doc. Dedent the rest of the doc after the first line so that it is formatted properly on the site.

PiperOrigin-RevId: 312564525
Change-Id: Icecb785b3914267411c3693d70b8daba3e06454d
This commit is contained in:
Yash Katariya 2020-05-20 15:14:01 -07:00 committed by TensorFlower Gardener
parent 33d5a54801
commit 4829f33e1f
2 changed files with 30 additions and 23 deletions

View File

@ -21,10 +21,11 @@ from __future__ import print_function
import contextlib
import numpy as np
import traceback
import weakref
import numpy as np
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@ -985,22 +986,20 @@ class TensorArray(object):
Example 3: A simple loop interacting with a `tf.Variable`.
# TODO(b/153898334) reenable this one flakyness is removed
# >>> v = tf.Variable(1)
# >>>
# >>> @tf.function
# ... def f(x):
# ... ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
# ...
# ... for i in tf.range(x):
# ... v.assign_add(i)
# ... ta = ta.write(i, v)
# ...
# ... return ta.stack()
# >>>
# >>> f(5)
# <tf.Tensor: shape=(5,), dtype=int32, numpy=array([ 1, 2, 4, 7, 11],
# dtype=int32)>
# TODO(b/153898334): Convert back to doctest once bug is resolved.
```
v = tf.Variable(1)
@tf.function
def f(x):
ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
for i in tf.range(x):
v.assign_add(i)
ta = ta.write(i, v)
return ta.stack()
f(5)
<tf.Tensor: shape=(5,), dtype=int32, numpy=array([ 1, 2, 4, 7, 11],
dtype=int32)>
```
"""
def __init__(self,

View File

@ -19,6 +19,7 @@ from __future__ import print_function
import copy
import sys
import textwrap
import traceback
import six # pylint: disable=unused-import
@ -231,20 +232,27 @@ def should_use_result(fn=None, warn_in_eager=False, error_in_function=False):
The wrapped function.
"""
def decorated(fn):
"""Decorates the input function."""
def wrapped(*args, **kwargs):
return _add_should_use_warning(fn(*args, **kwargs),
warn_in_eager=warn_in_eager,
error_in_function=error_in_function)
fn_doc = fn.__doc__ or ''
split_doc = fn_doc.split('\n', 1)
if len(split_doc) == 1:
updated_doc = fn_doc
else:
brief, rest = split_doc
updated_doc = '\n'.join([brief, textwrap.dedent(rest)])
note = ('\n\nNote: The output of this function should be used. If it is '
'not, a warning will be logged or an error may be raised. '
'To mark the output as used, call its .mark_used() method.')
return tf_decorator.make_decorator(
target=fn,
decorator_func=wrapped,
decorator_name='should_use_result',
decorator_doc=(
(fn.__doc__ or '') +
('\n\n '
'**NOTE** The output of this function should be used. If it is '
'not, a warning will be logged or an error may be raised. '
'To mark the output as used, call its .mark_used() method.')))
decorator_doc=updated_doc + note)
if fn is not None:
return decorated(fn)