Re-enable write-once, read-many semantics for TensorArray.

This implementation is a bit more efficient than the previous one because
the first write just performs a shallow copy.  Only on an aggregation is
any new memory allocated.

For read-many semantics, the operations read, pack, and concat must be called
with parameter  clear_after_read=False.  By default, the flag is set True; this
means a read will remove the reference to the underlying Tensor in
the TensorArray to reclaim memory in the runtime.
Change: 119404140

1 files changed, 143 insertions, 57 deletions

diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py
index 35f63db23f..b5bc6e2570 100644
--- a/tensorflow/python/ops/tensor_array_ops.py
+++ b/tensorflow/python/ops/tensor_array_ops.py
@@ -23,6 +23,7 @@ from __future__ import print_function
 from tensorflow.python.framework import dtypes as _dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import constant_op
 from tensorflow.python.ops import gen_data_flow_ops
 from tensorflow.python.ops import math_ops
@@ -53,7 +54,7 @@ class TensorArray(object):
   """
 
   def __init__(self, dtype, size=None, dynamic_size=None,
-               tensor_array_name=None,
+               clear_after_read=None, tensor_array_name=None,
                handle=None, flow=None, name=None):
     """Construct a new TensorArray or wrap an existing TensorArray handle.
 
@@ -63,6 +64,9 @@ class TensorArray(object):
         Required if handle is not provided.
       dynamic_size: (optional) Python bool: If true, writes to the TensorArray
         can grow the TensorArray past its initial size.  Default: False.
+      clear_after_read: Boolean (optional, default: True).  If True, clear
+        TensorArray values after reading them.  This disables read-many
+        semantics, but allows early release of memory.
       tensor_array_name: (optional) Python string: the name of the TensorArray.
         This is used when creating the TensorArray handle.  If this value is
         set, handle should be None.
@@ -89,7 +93,12 @@ class TensorArray(object):
     if handle is not None and dynamic_size is not None:
       raise ValueError("Cannot provide both a handle and dynamic_size "
                        "at the same time")
+    if handle is not None and clear_after_read is not None:
+      raise ValueError("Cannot provide both a handle and clear_after_read "
+                       "at the same time")
 
+    if clear_after_read is None:
+      clear_after_read = True
     dynamic_size = dynamic_size or False
 
     self._dtype = dtype
@@ -97,13 +106,22 @@ class TensorArray(object):
       if handle is not None:
         self._handle = handle
       else:
-        self._handle = gen_data_flow_ops._tensor_array(
-            dtype=dtype, size=size, dynamic_size=dynamic_size,
-            tensor_array_name=tensor_array_name, name=scope)
-    if flow is not None:
-      self._flow = flow
-    else:
-      self._flow = constant_op.constant(0, dtype=_dtypes.float32)
+        if flow is not None:
+          with ops.colocate_with(flow):
+            self._handle = gen_data_flow_ops._tensor_array(
+                dtype=dtype, size=size, dynamic_size=dynamic_size,
+                clear_after_read=clear_after_read,
+                tensor_array_name=tensor_array_name, name=scope)
+        else:
+          self._handle = gen_data_flow_ops._tensor_array(
+              dtype=dtype, size=size, dynamic_size=dynamic_size,
+              clear_after_read=clear_after_read,
+              tensor_array_name=tensor_array_name, name=scope)
+      if flow is not None:
+        self._flow = flow
+      else:
+        with ops.colocate_with(self._handle):
+          self._flow = constant_op.constant(0, dtype=_dtypes.float32)
 
   @property
   def flow(self):
@@ -120,80 +138,148 @@ class TensorArray(object):
     """The reference to the TensorArray."""
     return self._handle
 
-  def grad(self, source, flow=None):
+  def grad(self, source, flow=None, name=None):
     # tensor_array_grad requires a flow input when forward
     # TensorArrays are dynamically sized.  This forces the creation
     # of the grad TensorArray only once the final forward array's size
     # is fixed.
     if flow is None:
       flow = self.flow
-    g_handle = gen_data_flow_ops._tensor_array_grad(
-        handle=self._handle, source=source, flow_in=flow)
-    g = TensorArray(dtype=self._dtype, handle=g_handle, flow=flow)
-    return g
+    with ops.op_scope([self._handle], name, "TensorArrayGrad"):
+      with ops.colocate_with(self._handle):
+        g_handle = gen_data_flow_ops._tensor_array_grad(
+            handle=self._handle, source=source, flow_in=flow, name=name)
+        with ops.control_dependencies([g_handle]):
+          flow = array_ops.identity(flow, name="gradient_flow")
+        g = TensorArray(dtype=self._dtype, handle=g_handle, flow=flow)
+        return g
 
   def read(self, index, name=None):
-    """Read the value at location `index` in the TensorArray."""
-    value = gen_data_flow_ops._tensor_array_read(
-        handle=self._handle, index=index, flow_in=self._flow, dtype=self._dtype,
-        name=name)
-    return value
+    """Read the value at location `index` in the TensorArray.
+
+    Args:
+      index: 0-D.  int32 tensor with the index to read from.
+      name: A name for the operation (optional).
+
+    Returns:
+      The tensor at index `index`.
+    """
+    with ops.colocate_with(self._handle):
+      value = gen_data_flow_ops._tensor_array_read(
+          handle=self._handle, index=index, flow_in=self._flow,
+          dtype=self._dtype, name=name)
+      return value
 
   def write(self, index, value, name=None):
-    """Write `value` into index `index` of the TensorArray."""
-    flow_out = gen_data_flow_ops._tensor_array_write(
-        handle=self._handle, index=index, value=value, flow_in=self._flow,
-        name=name)
-    # Size below is ignored
-    ta = TensorArray(dtype=self._dtype, handle=self._handle)
-    ta._flow = flow_out
-    return ta
+    """Write `value` into index `index` of the TensorArray.
+
+    Args:
+      index: 0-D.  int32 scalar with the index to write to.
+      value: N-D.  Tensor of type `dtype`.  The Tensor to write to this index.
+      name: A name for the operation (optional).
+
+    Returns:
+      A new TensorArray object with flow that ensures the write occurs.
+      Use this object all for subsequent operations.
+    """
+    with ops.colocate_with(self._handle):
+      flow_out = gen_data_flow_ops._tensor_array_write(
+          handle=self._handle, index=index, value=value, flow_in=self._flow,
+          name=name)
+      ta = TensorArray(dtype=self._dtype, handle=self._handle)
+      ta._flow = flow_out
+      return ta
 
   def pack(self, name=None):
-    """Return the values in the TensorArray as a packed `Tensor`."""
-    value = gen_data_flow_ops._tensor_array_pack(
-        handle=self._handle, flow_in=self._flow, dtype=self._dtype,
-        name=name)
+    """Return the values in the TensorArray as a packed `Tensor`.
+
+    All of the values must have been written and their shapes must all match.
+
+    Args:
+      name: A name for the operation (optional).
+
+    Returns:
+      All the tensors in the TensorArray packed into one tensor.
+    """
+    with ops.colocate_with(self._handle):
+      value = gen_data_flow_ops._tensor_array_pack(
+          handle=self._handle, flow_in=self._flow, dtype=self._dtype,
+          name=name)
 
-    return value
+      return value
 
   def concat(self, name=None):
-    """Return the values in the TensorArray as a concatenated `Tensor`."""
-    value, _ = gen_data_flow_ops._tensor_array_concat(
-        handle=self._handle, flow_in=self._flow, dtype=self._dtype,
-        name=name)
-    return value
+    """Return the values in the TensorArray as a concatenated `Tensor`.
+
+    All of the values must have been written, their ranks must match, and
+    and their shapes must all match for all dimensions except the first.
+
+    Args:
+      name: A name for the operation (optional).
+
+    Returns:
+      All the tensors in the TensorArray concatenated into one tensor.
+    """
+    with ops.colocate_with(self._handle):
+      value, _ = gen_data_flow_ops._tensor_array_concat(
+          handle=self._handle, flow_in=self._flow, dtype=self._dtype,
+          name=name)
+      return value
 
   def unpack(self, value, name=None):
-    """Pack the values of a `Tensor` in the TensorArray."""
-    flow_out = gen_data_flow_ops._tensor_array_unpack(
-        handle=self._handle, value=value, flow_in=self._flow,
-        name=name)
-    ta = TensorArray(dtype=self._dtype, handle=self._handle)
-    ta._flow = flow_out
-    return ta
+    """Pack the values of a `Tensor` in the TensorArray.
+
+    Args:
+      value: (N+1)-D.  Tensor of type `dtype`.  The Tensor to unpack.
+      name: A name for the operation (optional).
+
+    Returns:
+      A new TensorArray object with flow that ensures the unpack occurs.
+      Use this object all for subsequent operations.
+    """
+    with ops.colocate_with(self._handle):
+      flow_out = gen_data_flow_ops._tensor_array_unpack(
+          handle=self._handle, value=value, flow_in=self._flow,
+          name=name)
+      ta = TensorArray(dtype=self._dtype, handle=self._handle)
+      ta._flow = flow_out
+      return ta
 
   def split(self, value, lengths, name=None):
-    """Split the values of a `Tensor` into the TensorArray."""
-    with ops.op_scope(
-        [self._handle, value, lengths], name, "TensorArraySplit"):
-      lengths = math_ops.to_int64(lengths)
-    flow_out = gen_data_flow_ops._tensor_array_split(
-        handle=self._handle, value=value, lengths=lengths, flow_in=self._flow,
-        name=name)
-    ta = TensorArray(dtype=self._dtype, handle=self._handle)
-    ta._flow = flow_out
-    return ta
+    """Split the values of a `Tensor` into the TensorArray.
+
+    Args:
+      value: (N+1)-D.  Tensor of type `dtype`.  The Tensor to split.
+      lengths: 1-D.  int32 vector with the lengths to use when splitting
+        `value` along its first dimension.
+      name: A name for the operation (optional).
+
+    Returns:
+      A new TensorArray object with flow that ensures the split occurs.
+      Use this object all for subsequent operations.
+    """
+    with ops.colocate_with(self._handle):
+      with ops.op_scope(
+          [self._handle, value, lengths], name, "TensorArraySplit"):
+        lengths = math_ops.to_int64(lengths)
+      flow_out = gen_data_flow_ops._tensor_array_split(
+          handle=self._handle, value=value, lengths=lengths, flow_in=self._flow,
+          name=name)
+      ta = TensorArray(dtype=self._dtype, handle=self._handle)
+      ta._flow = flow_out
+      return ta
 
   def size(self, name=None):
     """Return the size of the TensorArray."""
-    return gen_data_flow_ops._tensor_array_size(
-        handle=self._handle, flow_in=self.flow, name=name)
+    with ops.colocate_with(self._handle):
+      return gen_data_flow_ops._tensor_array_size(
+          handle=self._handle, flow_in=self.flow, name=name)
 
   def close(self, name=None):
     """Close the current TensorArray."""
-    return gen_data_flow_ops._tensor_array_close(
-        handle=self._handle, name=name)
+    with ops.colocate_with(self._handle):
+      return gen_data_flow_ops._tensor_array_close(
+          handle=self._handle, name=name)
 
 
 @ops.RegisterShape("TensorArray")