Enable doctest in TensorArray class.
The race condition happening earlier was because resource reads were not added to the list of control outputs of the tf.function. That has been fixed in a separate change.
PiperOrigin-RevId: 340880516
Change-Id: If8c23a1e6876611b1709c95bcfae0b95cce47c0b
diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py
index 58dc920..862f705 100644
--- a/tensorflow/python/ops/tensor_array_ops.py
+++ b/tensorflow/python/ops/tensor_array_ops.py
@@ -986,20 +986,17 @@
Example 3: A simple loop interacting with a `tf.Variable`.
- # 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)
+ >>> 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,
