Support a single Tensor in StagingArea.put() (tensorflow#17863)

imsheridan · sb2nov · commit f16b50e38854 · 2018-03-31T15:57:00.000-07:00
* Support a single Tensor in a StagingArea.put()

* Fix the input values parameter could also be dictionary type
diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py
@@ -1769,7 +1769,9 @@ def put(self, values, name=None):
     its capacity.
 
     Args:
-      values: Tensor (or a tuple of Tensors) to place into the staging area.
+      values: A single tensor, a list or tuple of tensors, or a dictionary with
+        tensor values. The number of elements must match the length of the
+        list provided to the dtypes argument when creating the StagingArea.
       name: A name for the operation (optional).
 
     Returns:
@@ -1780,11 +1782,12 @@ def put(self, values, name=None):
     """
     with ops.name_scope(name, "%s_put" % self._name,
                         self._scope_vals(values)) as scope:
+      
+      if not isinstance(values, (list, tuple, dict)):
+        values = [values]
 
       # Hard-code indices for this staging area
-      indices = (
-          list(six.moves.range(len(values)))
-          if isinstance(values, (list, tuple)) else None)
+      indices = list(six.moves.range(len(values)))
       vals, _ = self._check_put_dtypes(values, indices)
 
       with ops.colocate_with(self._coloc_op):
