"""Convert dtypes to a list of types."""

assert dtypes is not None

if not (isinstance(dtypes, list) or isinstance(dtypes, tuple)):

# We have a single type.

return [dtypes]

else:

# We have a list or tuple of types.

unknown_rank_allowed=False):

"""Convert shapes to a list of tuples of int (or None)."""

if unknown_dim_allowed:

if (not isinstance(shapes, collections.Sequence)

or not shapes

or any(shape is None or isinstance(shape, int) for shape in shapes)):

raise ValueError(

"When providing partial shapes, a list of shapes must be provided.")

if shapes is None: return None

if isinstance(shapes, tensor_shape.TensorShape):

shapes = [shapes]

if not isinstance(shapes, (tuple, list)):

raise TypeError(

"shapes must be a TensorShape or a list or tuple of TensorShapes.")

if all(shape is None or isinstance(shape, int) for shape in shapes):

# We have a single shape.

shapes = [shapes]

shapes = [tensor_shape.as_shape(shape) for shape in shapes]

if not unknown_dim_allowed:

if any([not shape.is_fully_defined() for shape in shapes]):

raise ValueError("All shapes must be fully defined: %s" % shapes)

if not unknown_rank_allowed:

if any([shape.dims is None for shape in shapes]):

raise ValueError("All shapes must have a defined rank: %s" % shapes)

if names is None:

return None

if not isinstance(names, (list, tuple)):

names = [names]

if len(names) != len(dtypes):

raise ValueError("List of names must have the same length as the list "

"of dtypes")

"""The greatest lower bound (ordered by specificity) TensorShape."""

s1 = tensor_shape.TensorShape(s1)

s2 = tensor_shape.TensorShape(s2)

if s1.ndims is None or s2.ndims is None or s1.ndims != s2.ndims:

return tensor_shape.unknown_shape()

d = [

d1 if d1 is not None and d1 == d2 else None

for (d1, d2) in zip(s1.as_list(), s2.as_list())]

"""Base class for queue implementations.

def __init__(self, dtypes, shapes, names, queue_ref):

"""Constructs a queue object from a queue reference.

self._dtypes = dtypes

if shapes is not None:

if len(shapes) != len(dtypes):

raise ValueError("Queue shapes must have the same length as dtypes")

self._shapes = [tensor_shape.TensorShape(s) for s in shapes]

else:

self._shapes = [tensor_shape.unknown_shape() for _ in self._dtypes]

if names is not None:

if len(names) != len(dtypes):

raise ValueError("Queue names must have the same length as dtypes")

self._names = names

else:

self._names = None

self._queue_ref = queue_ref

self._name = self._queue_ref.op.name.split("/")[-1]

@staticmethod

def from_list(index, queues):

"""Create a queue using the queue reference from `queues[index]`.

if ((not queues) or

(not isinstance(queues, list)) or

(not all(isinstance(x, QueueBase) for x in queues))):

raise TypeError("A list of queues expected")

dtypes = queues[0].dtypes

if not all([dtypes == q.dtypes for q in queues[1:]]):

raise TypeError("Queues do not have matching component dtypes.")

names = queues[0].names

if not all([names == q.names for q in queues[1:]]):

raise TypeError("Queues do not have matching component names.")

queue_shapes = [q.shapes for q in queues]

reduced_shapes = [

six.moves.reduce(_shape_common, s) for s in zip(*queue_shapes)]

queue_refs = array_ops.stack([x.queue_ref for x in queues])

selected_queue = array_ops.gather(queue_refs, index)

return QueueBase(dtypes=dtypes, shapes=reduced_shapes, names=names,

queue_ref=selected_queue)

@property

def queue_ref(self):

"""The underlying queue reference."""

return self._queue_ref

@property

def name(self):

"""The name of the underlying queue."""

return self._queue_ref.op.name

@property

def dtypes(self):

"""The list of dtypes for each component of a queue element."""

return self._dtypes

@property

def shapes(self):

"""The list of shapes for each component of a queue element."""

return self._shapes

@property

def names(self):

"""The list of names for each component of a queue element."""

return self._names

def _check_enqueue_dtypes(self, vals):

"""Validate and convert `vals` to a list of `Tensor`s.

if isinstance(vals, dict):

if not self._names:

raise ValueError("Queue must have names to enqueue a dictionary")

if sorted(self._names) != sorted(vals.keys()):

raise ValueError("Keys in dictionary to enqueue do not match "

"names of Queue. Dictionary: (%s), Queue: (%s)" %

(sorted(vals.keys()), sorted(self._names)))

# The order of values in `self._names` indicates the order in which the

# tensors in the dictionary `vals` must be listed.

vals = [vals[k] for k in self._names]

else:

if self._names:

raise ValueError("You must enqueue a dictionary in a Queue with names")

if not isinstance(vals, (list, tuple)):

vals = [vals]

tensors = []

for i, (val, dtype) in enumerate(zip(vals, self._dtypes)):

tensors.append(ops.convert_to_tensor(val, dtype=dtype,

name="component_%d" % i))

return tensors

def _scope_vals(self, vals):

"""Return a list of values to pass to `name_scope()`.

if isinstance(vals, (list, tuple)):

return vals

elif isinstance(vals, dict):

return vals.values()

else:

return [vals]

def enqueue(self, vals, name=None):

"""Enqueues one element to this queue.

with ops.name_scope(name, "%s_enqueue" % self._name,

self._scope_vals(vals)) as scope:

vals = self._check_enqueue_dtypes(vals)

# NOTE(mrry): Not using a shape function because we need access to

# the `QueueBase` object.

for val, shape in zip(vals, self._shapes):

val.get_shape().assert_is_compatible_with(shape)

if self._queue_ref.dtype == _dtypes.resource:

return gen_data_flow_ops._queue_enqueue_v2(

self._queue_ref, vals, name=scope)

else:

return gen_data_flow_ops._queue_enqueue(

self._queue_ref, vals, name=scope)

def enqueue_many(self, vals, name=None):

"""Enqueues zero or more elements to this queue.

with ops.name_scope(name, "%s_EnqueueMany" % self._name,

self._scope_vals(vals)) as scope:

vals = self._check_enqueue_dtypes(vals)

# NOTE(mrry): Not using a shape function because we need access to

# the `QueueBase` object.

batch_dim = vals[0].get_shape().with_rank_at_least(1)[0]

for val, shape in zip(vals, self._shapes):

batch_dim = batch_dim.merge_with(

val.get_shape().with_rank_at_least(1)[0])

val.get_shape()[1:].assert_is_compatible_with(shape)

return gen_data_flow_ops._queue_enqueue_many_v2(

self._queue_ref, vals, name=scope)

def _dequeue_return_value(self, tensors):

"""Return the value to return from a dequeue op.

if self._names:

# The returned values in `tensors` are in the same order as

# the names in `self._names`.

return {n: tensors[i] for i, n in enumerate(self._names)}

elif len(tensors) == 1:

return tensors[0]

else:

return tensors

def dequeue(self, name=None):

"""Dequeues one element from this queue.

if name is None:

name = "%s_Dequeue" % self._name

if self._queue_ref.dtype == _dtypes.resource:

ret = gen_data_flow_ops._queue_dequeue_v2(

self._queue_ref, self._dtypes, name=name)

else:

ret = gen_data_flow_ops._queue_dequeue(

self._queue_ref, self._dtypes, name=name)

# NOTE(mrry): Not using a shape function because we need access to

# the `QueueBase` object.

op = ret[0].op

for output, shape in zip(op.values(), self._shapes):

output.set_shape(shape)

return self._dequeue_return_value(ret)

def dequeue_many(self, n, name=None):

"""Dequeues and concatenates `n` elements from this queue.

if name is None:

name = "%s_DequeueMany" % self._name

ret = gen_data_flow_ops._queue_dequeue_many_v2(

self._queue_ref, n=n, component_types=self._dtypes, name=name)

# NOTE(mrry): Not using a shape function because we need access to

# the Queue object.

op = ret[0].op

batch_dim = tensor_shape.Dimension(tensor_util.constant_value(op.inputs[1]))

for output, shape in zip(op.values(), self._shapes):

output.set_shape(tensor_shape.TensorShape([batch_dim]).concatenate(shape))

return self._dequeue_return_value(ret)

def dequeue_up_to(self, n, name=None):

"""Dequeues and concatenates `n` elements from this queue.

if name is None:

name = "%s_DequeueUpTo" % self._name

ret = gen_data_flow_ops._queue_dequeue_up_to_v2(

self._queue_ref, n=n, component_types=self._dtypes, name=name)

# NOTE(mrry): Not using a shape function because we need access to

# the Queue object.

op = ret[0].op

for output, shape in zip(op.values(), self._shapes):

output.set_shape(tensor_shape.TensorShape([None]).concatenate(shape))

return self._dequeue_return_value(ret)

def close(self, cancel_pending_enqueues=False, name=None):

"""Closes this queue.

if name is None:

name = "%s_Close" % self._name

if self._queue_ref.dtype == _dtypes.resource:

return gen_data_flow_ops._queue_close_v2(

self._queue_ref, cancel_pending_enqueues=cancel_pending_enqueues,

name=name)

else:

return gen_data_flow_ops._queue_close(

self._queue_ref, cancel_pending_enqueues=cancel_pending_enqueues,

name=name)

def size(self, name=None):

"""Compute the number of elements in this queue.

if name is None:

name = "%s_Size" % self._name

if self._queue_ref.dtype == _dtypes.resource:

return gen_data_flow_ops._queue_size_v2(self._queue_ref, name=name)

else:

"""A queue implementation that dequeues elements in a random order.

def __init__(self, capacity, min_after_dequeue, dtypes, shapes=None,

names=None, seed=None, shared_name=None,

name="random_shuffle_queue"):

"""Create a queue that dequeues elements in a random order.

dtypes = _as_type_list(dtypes)

shapes = _as_shape_list(shapes, dtypes)

names = _as_name_list(names, dtypes)

seed1, seed2 = random_seed.get_seed(seed)

if seed1 is None and seed2 is None:

seed1, seed2 = 0, 0

elif seed is None and shared_name is not None:

# This means that graph seed is provided but op seed is not provided.

# If shared_name is also provided, make seed2 depend only on the graph

# seed and shared_name. (seed2 from get_seed() is generally dependent on

# the id of the last op created.)

string = (str(seed1) + shared_name).encode("utf-8")

seed2 = int(hashlib.md5(string).hexdigest()[:8], 16) & 0x7FFFFFFF

queue_ref = gen_data_flow_ops._random_shuffle_queue_v2(

component_types=dtypes, shapes=shapes, capacity=capacity,

min_after_dequeue=min_after_dequeue, seed=seed1, seed2=seed2,

shared_name=shared_name, name=name)

"""A queue implementation that dequeues elements in first-in first-out order.

def __init__(self, capacity, dtypes, shapes=None, names=None,

shared_name=None, name="fifo_queue"):

"""Creates a queue that dequeues elements in a first-in first-out order.

dtypes = _as_type_list(dtypes)

shapes = _as_shape_list(shapes, dtypes)

names = _as_name_list(names, dtypes)

queue_ref = gen_data_flow_ops._fifo_queue_v2(

component_types=dtypes, shapes=shapes, capacity=capacity,

shared_name=shared_name, name=name)

"""A FIFOQueue that supports batching variable-sized tensors by padding.

def __init__(self, capacity, dtypes, shapes, names=None, shared_name=None,

name="padding_fifo_queue"):

"""Creates a queue that dequeues elements in a first-in first-out order.

dtypes = _as_type_list(dtypes)

shapes = _as_shape_list(shapes, dtypes, unknown_dim_allowed=True)

names = _as_name_list(names, dtypes)

if len(dtypes) != len(shapes):

raise ValueError("Shapes must be provided for all components, "

"but received %d dtypes and %d shapes."

% (len(dtypes), len(shapes)))

queue_ref = gen_data_flow_ops._padding_fifo_queue_v2(

component_types=dtypes, shapes=shapes, capacity=capacity,

shared_name=shared_name, name=name)

"""A queue implementation that dequeues elements in prioritized order.

def __init__(self, capacity, types, shapes=None, names=None, shared_name=None,

name="priority_queue"):

"""Creates a queue that dequeues elements in a first-in first-out order.

types = _as_type_list(types)

shapes = _as_shape_list(shapes, types)

queue_ref = gen_data_flow_ops._priority_queue_v2(

component_types=types, shapes=shapes, capacity=capacity,

shared_name=shared_name, name=name)

priority_dtypes = [_dtypes.int64] + types

priority_shapes = [()] + shapes if shapes else shapes

super(PriorityQueue, self).__init__(

"""Represents a key-value map that persists across graph executions."""

def __init__(self, types, shapes=None, shared_name=None, name="barrier"):

"""Creates a barrier that persists across different graph executions.

self._types = _as_type_list(types)

if shapes is not None:

shapes = _as_shape_list(shapes, self._types)

self._shapes = [tensor_shape.TensorShape(s) for s in shapes]

for i, shape in enumerate(self._shapes):

if shape.num_elements() == 0:

raise ValueError("Empty tensors are not supported, but received "

"shape '%s' at index %d" % (shape, i))

else:

self._shapes = [tensor_shape.unknown_shape() for _ in self._types]

self._barrier_ref = gen_data_flow_ops._barrier(

component_types=self._types, shapes=self._shapes,

shared_name=shared_name, name=name)

self._name = self._barrier_ref.op.name.split("/")[-1]

@property

def barrier_ref(self):

"""Get the underlying barrier reference."""

return self._barrier_ref

@property

def name(self):

"""The name of the underlying barrier."""

return self._barrier_ref.op.name

def insert_many(self, component_index, keys, values, name=None):

"""For each key, assigns the respective value to the specified component.

if name is None:

name = "%s_BarrierInsertMany" % self._name

return gen_data_flow_ops._barrier_insert_many(

self._barrier_ref, keys, values, component_index, name=name)

def take_many(self,

num_elements,

allow_small_batch=False,

timeout=None,

name=None):

"""Takes the given number of completed elements from this barrier.

if name is None:

name = "%s_BarrierTakeMany" % self._name

ret = gen_data_flow_ops._barrier_take_many(self._barrier_ref,

num_elements,

self._types,

allow_small_batch,

timeout,

name=name)

# NOTE(mrry): Not using a shape function because we need access to

# the Barrier object.

op = ret[0].op

if allow_small_batch:

batch_dim = None

else:

batch_dim = tensor_shape.Dimension(

tensor_util.constant_value(op.inputs[1]))

op.outputs[0].set_shape(tensor_shape.vector(batch_dim)) # indices

op.outputs[1].set_shape(tensor_shape.vector(batch_dim)) # keys

for output, shape in zip(op.outputs[2:], self._shapes): # value_list

output.set_shape(tensor_shape.TensorShape([batch_dim]).concatenate(shape))

return ret

def close(self, cancel_pending_enqueues=False, name=None):

"""Closes this barrier.