Ejemplo n.º 1
0
 def testPartiallyDefinedShape(self):
     s = tensor_shape.TensorShape([
         tensor_shape.Dimension(3),
         tensor_shape.Dimension(None),
         tensor_shape.Dimension(7)
     ])
     # pylint: disable=g-error-prone-assert-raises
     with self.assertRaisesRegex(ValueError,
                                 "Shape .+ is not fully defined"):
         s.assert_is_fully_defined()
     # pylint: enable=g-error-prone-assert-raises
     self.assertEqual(s.rank, 3)
     self.assertLen(s, 3)
     self.assertTrue(s)
     s.assert_has_rank(3)
     self.assertEqual(tensor_shape.Dimension(3), s[0])
     self.assertEqual(tensor_shape.Dimension(None).value, s.dims[1].value)
     self.assertEqual(tensor_shape.Dimension(7), s.dims[2])
     s.assert_same_rank([6, 3, 7])
     for d1, d2 in zip(s, [3, None, 7]):
         assert tensor_shape.dimension_value(d1) == d2
Ejemplo n.º 2
0
def adaptive_embedding_lookup(params,
                              ids,
                              transforms,
                              max_norm=None,
                              name=None):
    """Looks up `ids` in a list of embedding tensors.
    This function is used to perform parallel lookups on the list of tensors in `params`.  It is a generalization of
    `tf.gather`, where `params` is interpreted as a partitioning of a large embedding tensor with different number
    of rows and columns. `params` may be a `PartitionedVariable`s as returned by using `tf.compat.v1.get_variable()`
    with a partitioner.
    Each element `id` of `ids` is partitioned between the elements of `params` according to the `div` partition
    strategy.
    Id space should be the same with total number of `params` rows.
    The results of the lookup are transformed with corresponding functions from `transforms` and concatenated into a
    dense tensor. The returned tensor shape selected by `transforms`.
    Args:
        params: A list of P tensors of different shape, representing sharded embedding tensors.  Alternatively, a
            `PartitionedVariable`, created by partitioning along dimension 0.
        ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked up in `params`.
        transforms: Functions applied to each retrieved embedding before concatenation. Required due to possible
            different embedding sizes of `params`.
        max_norm: If not `None`, each embedding is clipped if its l2-norm is larger than this value.
        name: A name for the operation (optional).
    Returns:
        A `Tensor` with the same type as the tensors in `params`.
    """

    if isinstance(ids, tf.RaggedTensor):
        return tf.ragged.map_flat_values(adaptive_embedding_lookup, params,
                                         ids, transforms, max_norm, name)

    if not isinstance(params, (list, tuple)) or len(params) < 2:
        raise ValueError('At least 2 variables required in params')

    if not isinstance(transforms,
                      (list, tuple)) or len(transforms) != len(params):
        raise ValueError('Each param should have corresponding transform')

    if not all([callable(t) for t in transforms]):
        raise ValueError('Each transform should be callable')

    with tf.name_scope(name or 'adaptive_embedding_lookup') as name:
        np = len(params)  # Number of partitions

        # Preserve the resource variable status to avoid accidental dense reads.
        if not any(
                isinstance(p, resource_variable_ops.ResourceVariable)
                for p in params):
            params = ops.convert_n_to_tensor_or_indexed_slices(params,
                                                               name='params')
        ids = tf.convert_to_tensor(ids, name='ids')

        # Flatten the ids. There is more than one params tensor.
        flat_ids = tf.reshape(ids, [-1])
        original_indices = tf.range(tf.size(flat_ids))

        # Create p_assignments and set new_ids for adaptive strategy.
        # Compute total_ids_capacity as the sum of dim-0 of params, then assign to
        # partitions based on a variable number of ids per partition. Optimize
        # if we already know the full shape statically.
        dim_0_sizes = []
        for p in range(np):
            param_p_dim = tensor_shape.Dimension(
                tensor_shape.dimension_value(params[p].get_shape()[0]))
            dim_0_sizes.append(param_p_dim)

        dim_0_size_value = sum(dim_0_sizes).value
        if dim_0_size_value:
            dim_0_sizes = tf.TensorShape(dim_0_sizes).as_list()
            total_ids_capacity = tf.constant(dim_0_size_value,
                                             dtype=flat_ids.dtype)
        else:
            dim_0_sizes = []
            for p in range(np):
                param_p_dim = tensor_shape.dimension_value(
                    params[p].get_shape()[0])
                if param_p_dim is not None:
                    dim_0_sizes.append(param_p_dim)
                else:
                    with ops.colocate_with(params[p]):
                        dim_0_sizes.append(tf.shape(params[p])[0])
            dim_0_sizes = tf.stack(dim_0_sizes)
            total_ids_capacity = tf.reduce_sum(dim_0_sizes)

        p_cumsum = tf.cumsum(tf.cast(dim_0_sizes, dtype=flat_ids.dtype))
        assert_max_id = tf.debugging.assert_less(
            tf.math.reduce_max(flat_ids), total_ids_capacity,
            'Invalid id. Maximum id should be less then total number of params rows'
        )
        with tf.control_dependencies([assert_max_id]):
            p_assignments = tf.searchsorted(
                p_cumsum,
                flat_ids,
                side='right',
            )

        # Cast partition assignments to int32 for use in dynamic_partition.
        # There really should not be more than 2^32 partitions.
        p_assignments = tf.cast(p_assignments, tf.int32)

        # Partition list of ids based on assignments into np separate lists
        p_intervals = tf.concat(([0], p_cumsum), 0)
        new_ids = flat_ids - tf.gather(p_intervals, p_assignments)
        gather_ids = tf.dynamic_partition(new_ids, p_assignments, np)

        # Similarly, partition the original indices.
        p_indices = tf.dynamic_partition(original_indices, p_assignments, np)

        # Do np separate lookups, finding embeddings for plist[p] in params[p]
        partitioned_result = []
        for p in range(np):
            pids = gather_ids[p]
            transform_fn = transforms[p]
            with ops.colocate_with(params[p]):
                result = tf.gather(params[p], pids)
                result = embedding_ops._clip(transform_fn(result), pids,
                                             max_norm)
            partitioned_result.append(result)

        # Stitch these back together
        ret = data_flow_ops.parallel_dynamic_stitch(p_indices,
                                                    partitioned_result,
                                                    name=name)

        # Determine the static element shape.
        element_shape_s = ret.get_shape()[1:]

        # Compute the dynamic element shape.
        if element_shape_s.is_fully_defined():
            element_shape_d = element_shape_s
        else:
            element_shape_d = tf.shape(ret)[1:]

        # Reshape to reverse the flattening of ids.
        ret = tf.reshape(ret, tf.concat([tf.shape(ids), element_shape_d], 0))

        # Normally the reshape is sufficient, but setting shape explicitly
        # teaches shape inference that params[1:].get_shape() matters.
        ret.set_shape(ids.get_shape().concatenate(element_shape_s))

        return ret
Ejemplo n.º 3
0
 def testConversion(self):
     """Make sure fully known TensorShape objects convert to Tensors."""
     shape = tensor_shape.TensorShape([1, tensor_shape.Dimension(2)])
     shape_tensor = tensor_util.shape_tensor(shape)
     self.assertAllEqual((1, 2), shape_tensor)
Ejemplo n.º 4
0
 def testUnknownDimension(self):
     dim = tensor_shape.Dimension(None)
     self.assertIs(None, dim.value)
     self.assertEqual(dim.value, tensor_shape.Dimension(None).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value,
         (dim + tensor_shape.Dimension(None)).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value,
         (dim * tensor_shape.Dimension(None)).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value,
         (dim // tensor_shape.Dimension(None)).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value,
         dim.merge_with(tensor_shape.Dimension(None)).value)
     self.assertIs(
         None,
         tensor_shape.Dimension(None) < tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) <= tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) > tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) >= tensor_shape.Dimension(None))
Ejemplo n.º 5
0
 def testDimension(self):
     dim = tensor_shape.Dimension(12)
     self.assertEqual(12, dim.value)
     self.assertEqual(12, int(dim))
     self.assertEqual(dim, tensor_shape.Dimension(12))
     self.assertEqual(tensor_shape.Dimension(15),
                      dim + tensor_shape.Dimension(3))
     self.assertEqual(tensor_shape.Dimension(15), dim + 3)
     self.assertEqual(tensor_shape.Dimension(24),
                      dim * tensor_shape.Dimension(2))
     self.assertEqual(tensor_shape.Dimension(24), dim * 2)
     self.assertEqual(tensor_shape.Dimension(6),
                      dim // tensor_shape.Dimension(2))
     self.assertEqual(tensor_shape.Dimension(6), dim // 2)
     self.assertEqual(tensor_shape.Dimension(12),
                      dim.merge_with(tensor_shape.Dimension(12)))
     self.assertEqual(tensor_shape.Dimension(12), dim.merge_with(12))
     self.assertLess(tensor_shape.Dimension(12), tensor_shape.Dimension(13))
     self.assertGreater(tensor_shape.Dimension(13),
                        tensor_shape.Dimension(12))
     self.assertLessEqual(tensor_shape.Dimension(12),
                          tensor_shape.Dimension(12))
     self.assertLessEqual(tensor_shape.Dimension(12),
                          tensor_shape.Dimension(13))
     self.assertGreater(tensor_shape.Dimension(13),
                        tensor_shape.Dimension(12))
     self.assertGreaterEqual(tensor_shape.Dimension(12),
                             tensor_shape.Dimension(12))
     self.assertGreaterEqual(tensor_shape.Dimension(13),
                             tensor_shape.Dimension(12))
     with self.assertRaises(ValueError):
         dim.merge_with(tensor_shape.Dimension(13))
Ejemplo n.º 6
0
    def take_many(self,
                  num_elements,
                  allow_small_batch=False,
                  timeout=None,
                  name=None):
        """Takes the given number of completed elements from this barrier.

    This operation concatenates completed-element component tensors along
    the 0th dimension to make a single component tensor.

    If barrier has no completed elements, this operation will block
    until there are 'num_elements' elements to take.

    Args:
      num_elements: The number of elements to take.
      allow_small_batch: If the barrier is closed, don't block if there are less
        completed elements than requested, but instead return all available
        completed elements.
        TODO(b/25743580): the semantics of `allow_small_batch` are experimental
        and may be extended to other cases in the future.
        TODO(ebrevdo): If a take_many(allow_small_batch=True) is blocking
        already when the barrier is closed, it will block for ever. Fix this
        by using asynchronous operations.
      timeout: This specifies the number of milliseconds to block
        before returning with DEADLINE_EXCEEDED. (This option is not
        supported yet.)
      name: A name for the operation (optional).

    Returns:
      A tuple of (index, key, value_list).
      "index" is a int64 tensor of length num_elements containing the
        index of the insert_many call for which the very first component of
        the given element was inserted into the Barrier, starting with
        the value -2**63.  Note, this value is different from the
        index of the insert_many call for which the element was completed.
      "key" is a string tensor of length num_elements containing the keys.
      "value_list" is a tuple of tensors, each one with size num_elements
        in the 0th dimension for each component in the barrier's values.

    """
        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 safe_embedding_lookup_sparse(
        embedding_weights,
        sparse_ids,
        sparse_weights=None,
        combiner="mean",
        default_id=None,
        name="safe_embedding_lookup_sparse",
        partition_strategy=None,  # no used
        max_norm=None,
        return_trainable=False):
    """ Provides a dynamic version of `tf.nn.safe_embedding_lookup_sparse`.

  Lookup embedding results, accounting for empty features and invalid weights.

  Any IDs will be treated as valid include non-positive IDs.
  Invalid weights (<= 0) are pruned from input weights, as well as any IDs
  with non-positive weight. For an entry with no features, the embedding vector
  for `default_id` is returned, or the 0-vector if `default_id` is not supplied.

  The ids and weights may be multi-dimensional. Embeddings are always aggregated
  along the last dimension.

  Args:
    embedding_weights: A single `dynamic_embedding.Variable` instance
      representing the complete embedding tensor.
    sparse_ids: `SparseTensor` of shape `[d_0, d_1, ..., d_n]` containing the
      ids. `d_0` is typically batch size.
    sparse_weights: `SparseTensor` of same shape as `sparse_ids`, containing
      float weights corresponding to `sparse_ids`, or `None` if all weights are
      be assumed to be 1.0.
    combiner: A string specifying how to combine embedding results for each
      entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean" the
      default.
    default_id: The id to use for an entry with no features.
    name: A name for this operation (optional).
    partition_strategy: A string specifying the partitioning strategy. Currently
      `"div"` and `"mod"` are supported. Default is `"div"`.
    max_norm: If not `None`, all embeddings are l2-normalized to max_norm before
      combining.

  Returns:
    combined_embeddings:
      A dense `Tensor` of shape `[d_0, d_1, ..., d_{n-1}, e_1, ..., e_m]`.
    trainable_wrap:
      A TrainableWrapper object used to fill the Optimizers `var_list`
        Only provided if `return_trainable` is True.

  Raises:
    ValueError: if `embedding_weights` is empty.
  """
    if embedding_weights is None:
        raise ValueError("Missing embedding_weights %s." % embedding_weights)

    if embedding_weights.key_dtype != sparse_ids.dtype:
        raise TypeError(
            "embedding_weights.key_dtype should be same with sparse_ids.dtype: "
            "{} vs. {}".format(embedding_weights.key_dtype, sparse_ids.dtype))

    weights_dtype = sparse_weights.dtype if sparse_weights is not None else None
    if weights_dtype and embedding_weights.value_dtype != weights_dtype:
        raise TypeError(
            "embedding_weights.value_dtype should be same with sparse_weights.dtype"
            ": {} vs. {}".format(embedding_weights.value_dtype, weights_dtype))

    scope = variable_scope.get_variable_scope()
    full_name = scope.name + "/" + name if scope.name else name
    with ops.name_scope(full_name + "/"):
        # Reshape higher-rank sparse ids and weights to linear segment ids.
        original_shape = sparse_ids.dense_shape
        original_rank_dim = tensor_shape.dimension_value(
            sparse_ids.dense_shape.get_shape()[0])
        original_rank = (array_ops.size(original_shape)
                         if original_rank_dim is None else original_rank_dim)
        sparse_ids = sparse_ops.sparse_reshape(sparse_ids, [
            math_ops.reduce_prod(
                array_ops.slice(original_shape, [0], [original_rank - 1])),
            array_ops.gather(original_shape, original_rank - 1)
        ])
        if sparse_weights is not None:
            sparse_weights = sparse_tensor.SparseTensor(
                sparse_ids.indices, sparse_weights.values,
                sparse_ids.dense_shape)

        # Prune invalid weights.
        if combiner != "sum":
            sparse_ids, sparse_weights = _prune_invalid_weights(
                sparse_ids, sparse_weights)

        # Fill in dummy values for empty features, if necessary.
        sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(
            sparse_ids, default_id or 0)
        if sparse_weights is not None:
            sparse_weights, _ = sparse_ops.sparse_fill_empty_rows(
                sparse_weights, 1.0)

        result, trainable_ = embedding_lookup_sparse(
            embedding_weights,
            sparse_ids,
            sparse_weights,
            combiner=combiner,
            partition_strategy=partition_strategy,
            name=name + "/embedding_lookup_sparse",
            max_norm=max_norm,
            return_trainable=True)

        if default_id is None:
            # Broadcast is_row_empty to the same shape as embedding_lookup_result,
            # for use in Select.
            is_row_empty = array_ops.tile(
                array_ops.reshape(is_row_empty, [-1, 1]),
                array_ops.stack([1, array_ops.shape(result)[1]]))

            result = array_ops.where(is_row_empty,
                                     array_ops.zeros_like(result),
                                     result,
                                     name="where")

        # Reshape back from linear ids back into higher-dimensional dense result.
        final_result = array_ops.reshape(
            result,
            array_ops.concat([
                array_ops.slice(math_ops.cast(original_shape, dtypes.int32),
                                [0], [original_rank - 1]),
                array_ops.slice(array_ops.shape(result), [1], [-1])
            ], 0))
        final_result.set_shape(
            tensor_shape.unknown_shape(
                (tensor_shape.Dimension(original_rank_dim) -
                 1).value).concatenate(result.get_shape()[1:]))
        return (final_result, trainable_) if return_trainable else final_result
Ejemplo n.º 8
0
 def testMergePartialShapes(self):
   s1 = tensor_shape.TensorShape([tensor_shape.Dimension(
       3), tensor_shape.Dimension(None), tensor_shape.Dimension(7)])
   s2 = tensor_shape.TensorShape([tensor_shape.Dimension(
       None), tensor_shape.Dimension(4), tensor_shape.Dimension(7)])
   self.assertEqual([3, 4, 7], s1.merge_with(s2).as_list())
Ejemplo n.º 9
0
 def testConcatenateWithDimension(self, concatenate_fn):
     tensor_shape.TensorShape([1, 2, 3]).assert_is_compatible_with(
         concatenate_fn(tensor_shape.TensorShape([1, 2]),
                        tensor_shape.Dimension(3)))
Ejemplo n.º 10
0
    def testMergeWithError(self):
        with self.assertRaisesRegex(TypeError, "must be integer or None"):
            tensor_shape.Dimension(42).merge_with([])

        with self.assertRaisesRegex(ValueError, "must be >= 0"):
            tensor_shape.Dimension(42).merge_with(-1)
Ejemplo n.º 11
0
    def testEquality(self):
        self.assertEqual(tensor_shape.Dimension(12),
                         tensor_shape.Dimension(12))
        self.assertNotEqual(tensor_shape.Dimension(12),
                            tensor_shape.Dimension(13))
        self.assertIsNone(
            tensor_shape.Dimension(12) == tensor_shape.Dimension(None))
        self.assertIsNone(
            tensor_shape.Dimension(None) == tensor_shape.Dimension(12))
        self.assertIsNone(
            tensor_shape.Dimension(None) == tensor_shape.Dimension(None))

        # None indicates ambiguous comparison, but comparison vs the wrong type
        # is unambiguously False.
        self.assertIsNotNone(tensor_shape.Dimension(None) == 12.99)
        self.assertNotEqual(tensor_shape.Dimension(None), 12.99)

        # pylint: disable=singleton-comparison, g-equals-none
        self.assertIsNone(tensor_shape.Dimension(None) == None)
        # pylint: enable=singleton-comparison, g-equals-none
        self.assertNotEqual(tensor_shape.Dimension(12), 12.99)
Ejemplo n.º 12
0
def map_fn(fn,
           elems,
           dtype=None,
           parallel_iterations=None,
           back_prop=True,
           swap_memory=False,
           infer_shape=True,
           name=None):
    """map on the list of tensors unpacked from `elems` on dimension 0.
  The simplest version of `map_fn` repeatedly applies the callable `fn` to a
  sequence of elements from first to last. The elements are made of the
  tensors unpacked from `elems`. `dtype` is the data type of the return
  value of `fn`. Users must provide `dtype` if it is different from
  the data type of `elems`.
  Suppose that `elems` is unpacked into `values`, a list of tensors. The shape
  of the result tensor is `[values.shape[0]] + fn(values[0]).shape`.
  This method also allows multi-arity `elems` and output of `fn`.  If `elems`
  is a (possibly nested) list or tuple of tensors, then each of these tensors
  must have a matching first (unpack) dimension.  The signature of `fn` may
  match the structure of `elems`.  That is, if `elems` is
  `(t1, [t2, t3, [t4, t5]])`, then an appropriate signature for `fn` is:
  `fn = lambda (t1, [t2, t3, [t4, t5]]):`.
  Furthermore, `fn` may emit a different structure than its input.  For example,
  `fn` may look like: `fn = lambda t1: return (t1 + 1, t1 - 1)`.  In this case,
  the `dtype` parameter is not optional: `dtype` must be a type or (possibly
  nested) tuple of types matching the output of `fn`.
  To apply a functional operation to the nonzero elements of a SparseTensor
  one of the following methods is recommended. First, if the function is
  expressible as TensorFlow ops, use
  ```python
    result = SparseTensor(input.indices, fn(input.values), input.dense_shape)
  ```
  If, however, the function is not expressible as a TensorFlow op, then use
  ```python
  result = SparseTensor(
    input.indices, map_fn(fn, input.values), input.dense_shape)
  ```
  instead.
  When executing eagerly, map_fn does not execute in parallel even if
  `parallel_iterations` is set to a value > 1. You can still get the
  performance benefits of running a function in parallel by using the
  `tf.contrib.eager.defun` decorator,
  ```python
  # Assume the function being used in map_fn is fn.
  # To ensure map_fn calls fn in parallel, use the defun decorator.
  @tf.contrib.eager.defun
  def func(tensor):
    return tf.map_fn(fn, tensor)
  ```
  Note that if you use the defun decorator, any non-TensorFlow Python code
  that you may have written in your function won't get executed. See
  `tf.contrib.eager.defun` for more details. The recommendation would be to
  debug without defun but switch to defun to get performance benefits of
  running map_fn in parallel.
  Args:
    fn: The callable to be performed.  It accepts one argument, which will
      have the same (possibly nested) structure as `elems`.  Its output
      must have the same structure as `dtype` if one is provided, otherwise
      it must have the same structure as `elems`.
    elems: A tensor or (possibly nested) sequence of tensors, each of which
      will be unpacked along their first dimension.  The nested sequence
      of the resulting slices will be applied to `fn`.
    dtype: (optional) The output type(s) of `fn`.  If `fn` returns a structure
      of Tensors differing from the structure of `elems`, then `dtype` is not
      optional and must have the same structure as the output of `fn`.
    parallel_iterations: (optional) The number of iterations allowed to run
      in parallel. When graph building, the default value is 10. While executing
      eagerly, the default value is set to 1.
    back_prop: (optional) True enables support for back propagation.
    swap_memory: (optional) True enables GPU-CPU memory swapping.
    infer_shape: (optional) False disables tests for consistent output shapes.
    name: (optional) Name prefix for the returned tensors.
  Returns:
    A tensor or (possibly nested) sequence of tensors.  Each tensor packs the
    results of applying `fn` to tensors unpacked from `elems` along the first
    dimension, from first to last.
  Raises:
    TypeError: if `fn` is not callable or the structure of the output of
      `fn` and `dtype` do not match, or if elems is a SparseTensor.
    ValueError: if the lengths of the output of `fn` and `dtype` do not match.
  Examples:
    ```python
    elems = np.array([1, 2, 3, 4, 5, 6])
    squares = map_fn(lambda x: x * x, elems)
    # squares == [1, 4, 9, 16, 25, 36]
    ```
    ```python
    elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
    alternate = map_fn(lambda x: x[0] * x[1], elems, dtype=tf.int64)
    # alternate == [-1, 2, -3]
    ```
    ```python
    elems = np.array([1, 2, 3])
    alternates = map_fn(lambda x: (x, -x), elems, dtype=(tf.int64, tf.int64))
    # alternates[0] == [1, 2, 3]
    # alternates[1] == [-1, -2, -3]
    ```
  """
    if not callable(fn):
        raise TypeError("fn must be callable.")

    if isinstance(elems, sparse_tensor.SparseTensor):
        raise TypeError(
            "To perform a map on the values of a sparse tensor use either "
            " SparseTensor(input.indices, fn(input.values), input.dense_shape) or "
            " SparseTensor(input.indices, map_fn(fn, input.values), "
            "input.dense_shape)")

    in_graph_mode = not context.executing_eagerly()
    # Set the default number of parallel_iterations depending on graph/eager mode.
    if in_graph_mode and not parallel_iterations:
        parallel_iterations = 10
    elif not in_graph_mode and not parallel_iterations:
        parallel_iterations = 1

    if not in_graph_mode and parallel_iterations > 1:
        logging.log_first_n(
            logging.WARN, "Setting parallel_iterations > 1 has no "
            "effect when executing eagerly. Consider calling map_fn"
            " with tf.contrib.eager.defun to execute fn in "
            "parallel.", 1)
        parallel_iterations = 1

    import pdb
    pdb.set_trace()
    input_is_sequence = nest.is_sequence(elems)
    '''
  input_flatten = lambda x: nest.flatten(x) if input_is_sequence else [x]
  def input_pack(x):
    return nest.pack_sequence_as(elems, x) if input_is_sequence else x[0]
  '''
    input_flatten = lambda x: x if input_is_sequence else [x]

    def input_pack(x):
        return x if input_is_sequence else x[0]

    if dtype is None:
        output_is_sequence = input_is_sequence
        output_flatten = input_flatten
        output_pack = input_pack
    else:
        output_is_sequence = nest.is_sequence(dtype)
        output_flatten = lambda x: nest.flatten(
            x) if output_is_sequence else [x]

        def output_pack(x):
            return (nest.pack_sequence_as(dtype, x)
                    if output_is_sequence else x[0])

    elems_flat = input_flatten(elems)

    with ops.name_scope(name, "map", elems_flat):
        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode:
            # Any get_variable calls in fn will cache the first call locally
            # and not issue repeated network I/O requests for each iteration.
            varscope = vs.get_variable_scope()
            varscope_caching_device_was_none = False
            if varscope.caching_device is None:
                # TODO(ebrevdo): Change to using colocate_with here and in other
                # methods.
                varscope.set_caching_device(lambda op: op.device)
                varscope_caching_device_was_none = True

        elems_flat = [
            ops.convert_to_tensor(elem, name="elem") for elem in elems_flat
        ]

        dtype = dtype or input_pack([elem.dtype for elem in elems_flat])
        dtype_flat = output_flatten(dtype)

        # Convert elems to tensor array. n may be known statically.
        static_shape = elems_flat[0].shape
        if static_shape.ndims is not None and static_shape.ndims < 1:
            if len(elems_flat) == 1:
                raise ValueError(
                    "elems must be a 1+ dimensional Tensor, not a scalar")
            else:
                raise ValueError(
                    "elements in elems must be 1+ dimensional Tensors, not scalars"
                )
        n = (tensor_shape.dimension_value(static_shape[0])
             or array_ops.shape(elems_flat[0])[0])

        # TensorArrays are always flat
        elems_ta = [
            tensor_array_ops.TensorArray(dtype=elem.dtype,
                                         size=n,
                                         dynamic_size=False,
                                         infer_shape=True)
            for elem in elems_flat
        ]
        # Unpack elements
        elems_ta = [
            elem_ta.unstack(elem)
            for elem_ta, elem in zip(elems_ta, elems_flat)
        ]

        i = constant_op.constant(0)

        accs_ta = [
            tensor_array_ops.TensorArray(dtype=dt,
                                         size=n,
                                         dynamic_size=False,
                                         infer_shape=infer_shape)
            for dt in dtype_flat
        ]

        def compute(i, tas):
            """The loop body of map_fn.
      Args:
        i: the loop counter
        tas: the flat TensorArray accumulator list
      Returns:
        (i + 1, tas): the updated counter + updated TensorArrays
      Raises:
        TypeError: if dtype and packed_fn_values structure do not match
        ValueType: if dtype and packed_fn_values lengths do not match
      """
            packed_values = input_pack(
                [elem_ta.read(i) for elem_ta in elems_ta])
            packed_fn_values = fn(packed_values)
            nest.assert_same_structure(dtype or elems, packed_fn_values)
            flat_fn_values = output_flatten(packed_fn_values)
            tas = [
                ta.write(i, value) for (ta, value) in zip(tas, flat_fn_values)
            ]
            return (i + 1, tas)

        _, r_a = control_flow_ops.while_loop(
            lambda i, _: i < n,
            compute, (i, accs_ta),
            parallel_iterations=parallel_iterations,
            back_prop=back_prop,
            swap_memory=swap_memory,
            maximum_iterations=n)
        results_flat = [r.stack() for r in r_a]

        n_static = tensor_shape.Dimension(
            tensor_shape.dimension_value(
                elems_flat[0].get_shape().with_rank_at_least(1)[0]))
        for elem in elems_flat[1:]:
            n_static.merge_with(
                tensor_shape.Dimension(
                    tensor_shape.dimension_value(
                        elem.get_shape().with_rank_at_least(1)[0])))
        for r in results_flat:
            r.set_shape(
                tensor_shape.TensorShape(n_static).concatenate(
                    r.get_shape()[1:]))

        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode and varscope_caching_device_was_none:
            varscope.set_caching_device(None)

        return output_pack(results_flat)
Ejemplo n.º 13
0
 def shape(self):
     feature_shape = self.tensors[0].shape[1:]
     batch_size = sum([tensor.shape[0] for tensor in self.tensors],
                      tensor_shape.Dimension(0))
     return tensor_shape.TensorShape([batch_size
                                      ]).concatenate(feature_shape)
Ejemplo n.º 14
0
 def output_shapes(self):
   num_elements = tensor_shape.Dimension(None)
   return (tensor_shape.matrix(num_elements, self._row_shape.shape[0] + 1),
           tensor_shape.vector(num_elements),
           tensor_shape.vector(self._row_shape.shape[0] + 1))
Ejemplo n.º 15
0
def constant_value_as_shape(tensor):  # pylint: disable=invalid-name
    """A version of `constant_value()` that returns a `TensorShape`.

  This version should be used when a constant tensor value is
  interpreted as a (possibly partial) shape, e.g. in the shape
  function for `tf.reshape()`. By explicitly requesting a
  `TensorShape` as the return value, it is possible to represent
  unknown dimensions; by contrast, `constant_value()` is
  all-or-nothing.

  Args:
    tensor: The rank-1 Tensor to be evaluated.

  Returns:
    A `TensorShape` based on the constant value of the given `tensor`.
  """
    shape = tensor.get_shape().with_rank(1)
    if tensor.get_shape() == [0]:
        return tensor_shape.scalar()
    elif tensor.op.type == "Shape":
        return tensor.op.inputs[0].get_shape()
    elif tensor.op.type == "Pack":
        ret = tensor_shape.scalar()  # Empty list.
        for pack_input in tensor.op.inputs:
            # `pack_input` must be a scalar. Attempt to evaluate it, and append it
            # to `ret`.
            pack_input_val = constant_value(pack_input)
            if pack_input_val is None or pack_input_val < 0:
                new_dim = tensor_shape.Dimension(None)
            else:
                new_dim = tensor_shape.Dimension(pack_input_val)
            ret = ret.concatenate([new_dim])
        return ret
    elif tensor.op.type == "Concat":
        # We assume that `tensor.op.inputs[0]` evaluates to 0, as this is
        # the only legal value when concatenating vectors, and it will
        # have been checked by a previous shape function.
        ret = tensor_shape.scalar()  # Empty list.
        for concat_input in tensor.op.inputs[1:]:
            # `concat_input` must be a vector. Attempt to evaluate it as a shape,
            # and concatenate it with `ret`.
            ret = ret.concatenate(constant_value_as_shape(concat_input))
        return ret
    elif tensor.op.type == "ConcatV2":
        # We assume that `tensor.op.inputs[-1]` evaluates to 0, as this is
        # the only legal value when concatenating vectors, and it will
        # have been checked by a previous shape function.
        ret = tensor_shape.scalar()  # Empty list.
        for concat_input in tensor.op.inputs[:-1]:
            # `concat_input` must be a vector. Attempt to evaluate it as a shape,
            # and concatenate it with `ret`.
            ret = ret.concatenate(constant_value_as_shape(concat_input))
        return ret
    else:
        ret = tensor_shape.unknown_shape(shape[0].value)
        value = constant_value(tensor)
        if value is not None:
            ret = ret.merge_with(
                tensor_shape.TensorShape(
                    [d if d != -1 else None for d in value]))
        return ret
Ejemplo n.º 16
0
 def testStr(self):
   self.assertEqual(str(tensor_shape.Dimension(7)), "7")
   self.assertEqual(str(tensor_shape.Dimension(None)), "?")
Ejemplo n.º 17
0
 def testUnsupportedType(self):
   with self.assertRaises(TypeError):
     tensor_shape.Dimension(dtypes.string)
Ejemplo n.º 18
0
 def test_axis_value(self):
     self.assertEqual(self.i_7.value, tensor_shape.Dimension(7))
     self.assertTrue(self.i_range.value == tuple(range(7)))
Ejemplo n.º 19
0
 def testDimension(self):
   dim = tensor_shape.Dimension(12)
   self.assertEqual(12, dim.value)
   self.assertEqual(12, int(dim))
   self.assertEqual(dim, tensor_shape.Dimension(12))
   self.assertEqual(tensor_shape.Dimension(15),
                    dim + tensor_shape.Dimension(3))
   self.assertEqual(tensor_shape.Dimension(15), dim + 3)
   self.assertEqual(tensor_shape.Dimension(15), 3 + dim)
   self.assertEqual(tensor_shape.Dimension(9), dim - 3)
   self.assertEqual(tensor_shape.Dimension(1), 13 - dim)
   self.assertEqual(tensor_shape.Dimension(24),
                    dim * tensor_shape.Dimension(2))
   self.assertEqual(tensor_shape.Dimension(24), dim * 2)
   self.assertEqual(tensor_shape.Dimension(24), 2 * dim)
   self.assertEqual([4] * 12, [4] * dim)
   self.assertEqual(12 * [4], dim * [4])
   self.assertEqual(tensor_shape.Dimension(24), 2 * dim)
   self.assertEqual(
       tensor_shape.Dimension(6), dim // tensor_shape.Dimension(2))
   self.assertEqual(tensor_shape.Dimension(6), dim // 2)
   self.assertEqual(tensor_shape.Dimension(0), 2 // dim)
   self.assertEqual(tensor_shape.Dimension(12),
                    dim.merge_with(tensor_shape.Dimension(12)))
   self.assertEqual(tensor_shape.Dimension(12), dim.merge_with(12))
   self.assertLess(tensor_shape.Dimension(12), tensor_shape.Dimension(13))
   self.assertGreater(tensor_shape.Dimension(13), tensor_shape.Dimension(12))
   self.assertLessEqual(tensor_shape.Dimension(12), tensor_shape.Dimension(12))
   self.assertLessEqual(tensor_shape.Dimension(12), tensor_shape.Dimension(13))
   self.assertGreater(tensor_shape.Dimension(13), tensor_shape.Dimension(12))
   self.assertGreaterEqual(tensor_shape.Dimension(12),
                           tensor_shape.Dimension(12))
   self.assertGreaterEqual(tensor_shape.Dimension(13),
                           tensor_shape.Dimension(12))
   self.assertNotEqual(dim, (12,))
   with self.assertRaises(ValueError):
     dim.merge_with(tensor_shape.Dimension(13))
Ejemplo n.º 20
0
def safe_embedding_lookup_sparse(embedding_weights,
                                 sparse_ids,
                                 sparse_weights=None,
                                 combiner="mean",
                                 default_id=None,
                                 name=None,
                                 partition_strategy="div",
                                 max_norm=None):
    """Lookup embedding results, accounting for invalid IDs and empty features.

  The partitioned embedding in `embedding_weights` must all be the same shape
  except for the first dimension. The first dimension is allowed to vary as the
  vocabulary size is not necessarily a multiple of `P`.  `embedding_weights`
  may be a `PartitionedVariable` as returned by using
  `tf.compat.v1.get_variable()` with a
  partitioner.

  Invalid IDs (< 0) are pruned from input IDs and weights, as well as any IDs
  with non-positive weight. For an entry with no features, the embedding vector
  for `default_id` is returned, or the 0-vector if `default_id` is not supplied.

  The ids and weights may be multi-dimensional. Embeddings are always aggregated
  along the last dimension.

  Args:
    embedding_weights:  A list of `P` float `Tensor`s or values representing
      partitioned embedding `Tensor`s.  Alternatively, a `PartitionedVariable`
      created by partitioning along dimension 0.  The total unpartitioned shape
      should be `[e_0, e_1, ..., e_m]`, where `e_0` represents the vocab size
      and `e_1, ..., e_m` are the embedding dimensions.
    sparse_ids: `SparseTensor` of shape `[d_0, d_1, ..., d_n]` containing the
      ids. `d_0` is typically batch size.
    sparse_weights: `SparseTensor` of same shape as `sparse_ids`, containing
      float weights corresponding to `sparse_ids`, or `None` if all weights are
      be assumed to be 1.0.
    combiner: A string specifying how to combine embedding results for each
      entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean" the
      default.
    default_id: The id to use for an entry with no features.
    name: A name for this operation (optional).
    partition_strategy: A string specifying the partitioning strategy. Currently
      `"div"` and `"mod"` are supported. Default is `"div"`.
    max_norm: If not `None`, all embeddings are l2-normalized to max_norm before
      combining.

  Returns:
    Dense `Tensor` of shape `[d_0, d_1, ..., d_{n-1}, e_1, ..., e_m]`.

  Raises:
    ValueError: if `embedding_weights` is empty.
  """
    if embedding_weights is None:
        raise ValueError("Missing embedding_weights %s." % embedding_weights)
    if isinstance(embedding_weights, variables.PartitionedVariable):
        embedding_weights = list(
            embedding_weights)  # get underlying Variables.
    if not isinstance(embedding_weights, list):
        embedding_weights = [embedding_weights]
    if len(embedding_weights) < 1:
        raise ValueError("Missing embedding_weights %s." % embedding_weights)

    dtype = sparse_weights.dtype if sparse_weights is not None else None
    embedding_weights = [
        w if (isinstance(w, resource_variable_ops.ResourceVariable)
              and dtype in (None, w.dtype)) else ops.convert_to_tensor(
                  w, dtype=dtype) for w in embedding_weights
    ]

    with ops.name_scope(name, "embedding_lookup", embedding_weights +
                        [sparse_ids, sparse_weights]) as scope:
        # Reshape higher-rank sparse ids and weights to linear segment ids.
        original_shape = sparse_ids.dense_shape
        original_rank_dim = tensor_shape.dimension_value(
            sparse_ids.dense_shape.get_shape()[0])
        original_rank = (array_ops.size(original_shape)
                         if original_rank_dim is None else original_rank_dim)
        sparse_ids = sparse_ops.sparse_reshape(sparse_ids, [
            math_ops.reduce_prod(
                array_ops.slice(original_shape, [0], [original_rank - 1])),
            array_ops.gather(original_shape, original_rank - 1)
        ])
        if sparse_weights is not None:
            sparse_weights = sparse_tensor.SparseTensor(
                sparse_ids.indices, sparse_weights.values,
                sparse_ids.dense_shape)

        # Prune invalid ids and weights.
        sparse_ids, sparse_weights = _prune_invalid_ids(
            sparse_ids, sparse_weights)
        if combiner != "sum":
            sparse_ids, sparse_weights = _prune_invalid_weights(
                sparse_ids, sparse_weights)

        # Fill in dummy values for empty features, if necessary.
        sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(
            sparse_ids, default_id or 0)
        if sparse_weights is not None:
            sparse_weights, _ = sparse_ops.sparse_fill_empty_rows(
                sparse_weights, 1.0)

        result = embedding_lookup_sparse(
            embedding_weights,
            sparse_ids,
            sparse_weights,
            combiner=combiner,
            partition_strategy=partition_strategy,
            name=None if default_id is None else scope,
            max_norm=max_norm)

        if default_id is None:
            # Broadcast is_row_empty to the same shape as embedding_lookup_result,
            # for use in Select.
            is_row_empty = array_ops.tile(
                array_ops.reshape(is_row_empty, [-1, 1]),
                array_ops.stack([1, array_ops.shape(result)[1]]))

            result = array_ops.where(is_row_empty,
                                     array_ops.zeros_like(result),
                                     result,
                                     name=scope)

        # Reshape back from linear ids back into higher-dimensional dense result.
        final_result = array_ops.reshape(
            result,
            array_ops.concat([
                array_ops.slice(math_ops.cast(original_shape, dtypes.int32),
                                [0], [original_rank - 1]),
                array_ops.slice(array_ops.shape(result), [1], [-1])
            ], 0))
        final_result.set_shape(
            tensor_shape.unknown_shape(
                (tensor_shape.Dimension(original_rank_dim) -
                 1).value).concatenate(result.get_shape()[1:]))
        return final_result
Ejemplo n.º 21
0
def scan(fn,
         elems,
         initializer=None,
         parallel_iterations=10,
         back_prop=True,
         swap_memory=False,
         infer_shape=True,
         reverse=False,
         name=None):
    """scan on the list of tensors unpacked from `elems` on dimension 0.

  The simplest version of `scan` repeatedly applies the callable `fn` to a
  sequence of elements from first to last. The elements are made of the tensors
  unpacked from `elems` on dimension 0. The callable fn takes two tensors as
  arguments. The first argument is the accumulated value computed from the
  preceding invocation of fn, and the second is the value at the current
  position of `elems`. If `initializer` is None, `elems` must contain at least
  one element, and its first element is used as the initializer.

  Suppose that `elems` is unpacked into `values`, a list of tensors. The shape
  of the result tensor is `[len(values)] + fn(initializer, values[0]).shape`.
  If reverse=True, it's fn(initializer, values[-1]).shape.

  This method also allows multi-arity `elems` and accumulator.  If `elems`
  is a (possibly nested) list or tuple of tensors, then each of these tensors
  must have a matching first (unpack) dimension.  The second argument of
  `fn` must match the structure of `elems`.

  If no `initializer` is provided, the output structure and dtypes of `fn`
  are assumed to be the same as its input; and in this case, the first
  argument of `fn` must match the structure of `elems`.

  If an `initializer` is provided, then the output of `fn` must have the same
  structure as `initializer`; and the first argument of `fn` must match
  this structure.

  For example, if `elems` is `(t1, [t2, t3])` and `initializer` is
  `[i1, i2]` then an appropriate signature for `fn` in `python2` is:
  `fn = lambda (acc_p1, acc_p2), (t1, [t2, t3]):` and `fn` must return a list,
  `[acc_n1, acc_n2]`.  An alternative correct signature for `fn`, and the
   one that works in `python3`, is:
  `fn = lambda a, t:`, where `a` and `t` correspond to the input tuples.

  Args:
    fn: The callable to be performed.  It accepts two arguments.  The first will
      have the same structure as `initializer` if one is provided, otherwise it
      will have the same structure as `elems`.  The second will have the same
      (possibly nested) structure as `elems`.  Its output must have the same
      structure as `initializer` if one is provided, otherwise it must have the
      same structure as `elems`.
    elems: A tensor or (possibly nested) sequence of tensors, each of which will
      be unpacked along their first dimension.  The nested sequence of the
      resulting slices will be the first argument to `fn`.
    initializer: (optional) A tensor or (possibly nested) sequence of tensors,
      initial value for the accumulator, and the expected output type of `fn`.
    parallel_iterations: (optional) The number of iterations allowed to run in
      parallel.
    back_prop: (optional) True enables support for back propagation.
    swap_memory: (optional) True enables GPU-CPU memory swapping.
    infer_shape: (optional) False disables tests for consistent output shapes.
    reverse: (optional) True scans the tensor last to first (instead of first to
      last).
    name: (optional) Name prefix for the returned tensors.

  Returns:
    A tensor or (possibly nested) sequence of tensors.  Each tensor packs the
    results of applying `fn` to tensors unpacked from `elems` along the first
    dimension, and the previous accumulator value(s), from first to last (or
    last to first, if `reverse=True`).

  Raises:
    TypeError: if `fn` is not callable or the structure of the output of
      `fn` and `initializer` do not match.
    ValueError: if the lengths of the output of `fn` and `initializer`
      do not match.

  Examples:
    ```python
    elems = np.array([1, 2, 3, 4, 5, 6])
    sum = scan(lambda a, x: a + x, elems)
    # sum == [1, 3, 6, 10, 15, 21]
    sum = scan(lambda a, x: a + x, elems, reverse=True)
    # sum == [21, 20, 18, 15, 11, 6]
    ```

    ```python
    elems = np.array([1, 2, 3, 4, 5, 6])
    initializer = np.array(0)
    sum_one = scan(
        lambda a, x: x[0] - x[1] + a, (elems + 1, elems), initializer)
    # sum_one == [1, 2, 3, 4, 5, 6]
    ```

    ```python
    elems = np.array([1, 0, 0, 0, 0, 0])
    initializer = (np.array(0), np.array(1))
    fibonaccis = scan(lambda a, _: (a[1], a[0] + a[1]), elems, initializer)
    # fibonaccis == ([1, 1, 2, 3, 5, 8], [1, 2, 3, 5, 8, 13])
    ```
  """
    if not callable(fn):
        raise TypeError("fn must be callable.")

    input_is_sequence = nest.is_sequence(elems)
    input_flatten = lambda x: nest.flatten(x) if input_is_sequence else [x]

    def input_pack(x):
        return nest.pack_sequence_as(elems, x) if input_is_sequence else x[0]

    if initializer is None:
        output_is_sequence = input_is_sequence
        output_flatten = input_flatten
        output_pack = input_pack
    else:
        output_is_sequence = nest.is_sequence(initializer)
        output_flatten = lambda x: nest.flatten(
            x) if output_is_sequence else [x]

        def output_pack(x):
            return (nest.pack_sequence_as(initializer, x)
                    if output_is_sequence else x[0])

    elems_flat = input_flatten(elems)

    in_graph_mode = not context.executing_eagerly()
    with ops.name_scope(name, "scan", elems_flat):
        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode:
            # Any get_variable calls in fn will cache the first call locally
            # and not issue repeated network I/O requests for each iteration.
            varscope = vs.get_variable_scope()
            varscope_caching_device_was_none = False
            if varscope.caching_device is None:
                # TODO(ebrevdo): Change to using colocate_with here and in other
                # methods.
                varscope.set_caching_device(lambda op: op.device)
                varscope_caching_device_was_none = True

        # Convert elems to tensor array.
        elems_flat = [
            ops.convert_to_tensor(elem, name="elem") for elem in elems_flat
        ]

        # Convert elems to tensor array. n may be known statically.
        n = tensor_shape.dimension_value(elems_flat[0].shape[0])
        if n is None:
            n = array_ops.shape(elems_flat[0])[0]

        # TensorArrays are always flat
        elems_ta = [
            tensor_array_ops.TensorArray(dtype=elem.dtype,
                                         size=n,
                                         dynamic_size=False,
                                         element_shape=elem.shape[1:],
                                         infer_shape=True)
            for elem in elems_flat
        ]
        # Unpack elements
        elems_ta = [
            elem_ta.unstack(elem)
            for elem_ta, elem in zip(elems_ta, elems_flat)
        ]

        if initializer is None:
            a_flat = [elem.read(n - 1 if reverse else 0) for elem in elems_ta]
            i = 1
        else:
            initializer_flat = output_flatten(initializer)
            a_flat = [ops.convert_to_tensor(init) for init in initializer_flat]
            i = 0

        # Create a tensor array to store the intermediate values.
        accs_ta = [
            tensor_array_ops.TensorArray(
                dtype=init.dtype,
                size=n,
                element_shape=init.shape if infer_shape else None,
                dynamic_size=False,
                infer_shape=infer_shape) for init in a_flat
        ]

        if initializer is None:
            accs_ta = [
                acc_ta.write(n - 1 if reverse else 0, a)
                for (acc_ta, a) in zip(accs_ta, a_flat)
            ]

        def compute(i, a_flat, tas):
            """The loop body of scan.

      Args:
        i: the loop counter.
        a_flat: the accumulator value(s), flattened.
        tas: the output accumulator TensorArray(s), flattened.

      Returns:
        [i + 1, a_flat, tas]: the updated counter + new accumulator values +
          updated TensorArrays

      Raises:
        TypeError: if initializer and fn() output structure do not match
        ValueType: if initializer and fn() output lengths do not match
      """
            packed_elems = input_pack(
                [elem_ta.read(i) for elem_ta in elems_ta])
            packed_a = output_pack(a_flat)
            a_out = fn(packed_a, packed_elems)
            nest.assert_same_structure(
                elems if initializer is None else initializer, a_out)
            flat_a_out = output_flatten(a_out)
            tas = [ta.write(i, value) for (ta, value) in zip(tas, flat_a_out)]
            if reverse:
                next_i = i - 1
            else:
                next_i = i + 1
            return (next_i, flat_a_out, tas)

        if reverse:
            initial_i = n - 1 - i
            condition = lambda i, _1, _2: i >= 0
        else:
            initial_i = i
            condition = lambda i, _1, _2: i < n
        _, _, r_a = control_flow_ops.while_loop(
            condition,
            compute, (initial_i, a_flat, accs_ta),
            parallel_iterations=parallel_iterations,
            back_prop=back_prop,
            swap_memory=swap_memory,
            maximum_iterations=n)

        results_flat = [r.stack() for r in r_a]

        n_static = tensor_shape.Dimension(
            tensor_shape.dimension_value(
                elems_flat[0].get_shape().with_rank_at_least(1)[0]))
        for elem in elems_flat[1:]:
            n_static.merge_with(
                tensor_shape.Dimension(
                    tensor_shape.dimension_value(
                        elem.get_shape().with_rank_at_least(1)[0])))
        for r in results_flat:
            r.set_shape(
                tensor_shape.TensorShape(n_static).concatenate(
                    r.get_shape()[1:]))

        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode and varscope_caching_device_was_none:
            varscope.set_caching_device(None)

        return output_pack(results_flat)
Ejemplo n.º 22
0
def _embedding_lookup_and_transform(params,
                                    ids,
                                    partition_strategy="mod",
                                    name=None,
                                    max_norm=None,
                                    transform_fn=None):
    """Helper function for embedding_lookup and _compute_sampled_logits.

  This function is a generalization of embedding_lookup that optionally
  applies a caller-specified transformation to each embedding. This is
  done through the `transform_fn` argument. If provided, the function is
  applied to each partitioned tensor of retrieved embeddings, colocated
  with the embeddings. This function will be called with a single `Tensor`
  argument of the same type as the `params` tensor and should return a
  `Tensor`. The shape of the argument will be the same as `params` except
  for the size of the first dimension. The first dimension of the result's
  shape must be the same size as the argument's.

  Args:
    params: See embedding_lookup.
    ids: See embedding_lookup.
    partition_strategy: See embedding_lookup.
    name: See embedding_lookup.
    max_norm: See embedding_lookup.
    transform_fn: An optional function to apply to each retrieved embedding. If
      max_norm is provided, transform_fn is applied to the norm-limited
      embeddings.

  Returns:
    See embedding_lookup for details.
  Raises:
    ValueError: If `params` is empty.
  """
    if params is None:
        raise ValueError("params must be specified")
    if isinstance(params, (list, tuple)) and not params:
        raise ValueError("Need at least one param")
    if isinstance(params, variables.PartitionedVariable):
        params = list(params)  # Iterate to get the underlying Variables.
    if not isinstance(params, list):
        params = [params]

    with ops.name_scope(name, "embedding_lookup", params + [ids]) as name:
        np = len(params)  # Number of partitions
        # Preserve the resource variable status to avoid accidental dense reads.
        if not any(
                isinstance(p, resource_variable_ops.ResourceVariable)
                for p in params):
            params = ops.convert_n_to_tensor_or_indexed_slices(params,
                                                               name="params")
        ids = ops.convert_to_tensor(ids, name="ids")
        if np == 1 and (not transform_fn or ids.get_shape().ndims == 1):
            with ops.colocate_with(params[0]):
                result = _clip(array_ops.gather(params[0], ids, name=name),
                               ids, max_norm)
                if transform_fn:
                    result = transform_fn(result)
            # Make sure the final result does not have colocation contraints on the
            # params. Similar to the case np > 1 where parallel_dynamic_stitch is
            # outside the scioe of all with ops.colocate_with(params[p]).
            return array_ops.identity(result)
        else:
            # Flatten the ids. There are two cases where we need to do this.
            # - There is more than one params tensor.
            # - There is a transform_fn and ids is not statically known to be 1-D.
            #   We must flatten in this case because transform_fn expects a flat
            #   tensor of embeddings.
            flat_ids = array_ops.reshape(ids, [-1])
            original_indices = math_ops.range(array_ops.size(flat_ids))

            # Create p_assignments and set new_ids depending on the strategy.
            if partition_strategy == "mod":
                p_assignments = flat_ids % np
                new_ids = flat_ids // np
            elif partition_strategy == "div":
                # Compute num_total_ids as the sum of dim-0 of params, then assign to
                # partitions based on a constant number of ids per partition. Optimize
                # if we already know the full shape statically.
                dim_0_size = tensor_shape.Dimension(
                    tensor_shape.dimension_value(params[0].get_shape()[0]))
                for p in xrange(1, np):
                    dim_0_size += tensor_shape.Dimension(
                        tensor_shape.dimension_value(params[p].get_shape()[0]))
                if dim_0_size.value:
                    num_total_ids = constant_op.constant(
                        dim_0_size.value, flat_ids.dtype)
                else:
                    dim_0_sizes = []
                    for p in xrange(np):
                        param_p_dim = tensor_shape.dimension_value(
                            params[p].get_shape()[0])
                        if param_p_dim is not None:
                            dim_0_sizes.append(param_p_dim)
                        else:
                            with ops.colocate_with(params[p]):
                                dim_0_sizes.append(
                                    array_ops.shape(params[p])[0])
                    num_total_ids = math_ops.reduce_sum(
                        math_ops.cast(array_ops.stack(dim_0_sizes),
                                      flat_ids.dtype))
                ids_per_partition = num_total_ids // np
                extras = num_total_ids % np

                p_assignments = math_ops.maximum(
                    flat_ids // (ids_per_partition + 1),
                    (flat_ids - extras) // ids_per_partition)

                # Emulate a conditional using a boolean indicator tensor
                new_ids = array_ops.where(
                    p_assignments < extras, flat_ids % (ids_per_partition + 1),
                    (flat_ids - extras) % ids_per_partition)
            else:
                raise ValueError("Unrecognized partition strategy: " +
                                 partition_strategy)

            # Cast partition assignments to int32 for use in dynamic_partition.
            # There really should not be more than 2^32 partitions.
            p_assignments = math_ops.cast(p_assignments, dtypes.int32)
            # Partition list of ids based on assignments into np separate lists
            gather_ids = data_flow_ops.dynamic_partition(
                new_ids, p_assignments, np)
            # Similarly, partition the original indices.
            pindices = data_flow_ops.dynamic_partition(original_indices,
                                                       p_assignments, np)
            # Do np separate lookups, finding embeddings for plist[p] in params[p]
            partitioned_result = []
            for p in xrange(np):
                pids = gather_ids[p]
                with ops.colocate_with(params[p]):
                    result = array_ops.gather(params[p], pids)
                    if transform_fn:
                        # If transform_fn is provided, the clip_by_norm precedes
                        # the transform and hence must be co-located. See below
                        # for the counterpart if transform_fn is not proveded.
                        result = transform_fn(_clip(result, pids, max_norm))
                partitioned_result.append(result)
            # Stitch these back together
            ret = data_flow_ops.parallel_dynamic_stitch(pindices,
                                                        partitioned_result,
                                                        name=name)

            # Determine the static element shape.
            if transform_fn is None:
                element_shape_s = params[0].get_shape()[1:]
                for p in params[1:]:
                    element_shape_s = element_shape_s.merge_with(
                        p.get_shape()[1:])
            else:
                element_shape_s = ret.get_shape()[1:]

            # Compute the dynamic element shape.
            if element_shape_s.is_fully_defined():
                element_shape_d = element_shape_s
            elif transform_fn is None:
                # It's important that we compute params[0].shape on the right device
                # to avoid data motion.
                with ops.colocate_with(params[0]):
                    params_shape = array_ops.shape(params[0])
                element_shape_d = params_shape[1:]
            else:
                element_shape_d = array_ops.shape(ret)[1:]

            # Reshape to reverse the flattening of ids.
            ret = array_ops.reshape(
                ret,
                array_ops.concat([array_ops.shape(ids), element_shape_d], 0))

            # Normally the reshape is sufficient, but setting shape explicitly
            # teaches shape inference that params[1:].get_shape() matters
            # (in the case that transform_fn is None).
            ret.set_shape(ids.get_shape().concatenate(element_shape_s))
            if not transform_fn:
                # If transform_fn was provided, the clip_by_norm was done above.
                ret = _clip(ret, ids, max_norm)
            return ret
Ejemplo n.º 23
0
def constant_value_as_shape(tensor):  # pylint: disable=invalid-name
  """A version of `constant_value()` that returns a `TensorShape`.

  This version should be used when a constant tensor value is
  interpreted as a (possibly partial) shape, e.g. in the shape
  function for `tf.reshape()`. By explicitly requesting a
  `TensorShape` as the return value, it is possible to represent
  unknown dimensions; by contrast, `constant_value()` is
  all-or-nothing.

  Args:
    tensor: The rank-0 or rank-1 Tensor to be evaluated.

  Returns:
    A `TensorShape` based on the constant value of the given `tensor`.

  Raises:
    ValueError: If the shape is rank-0 and is not statically known to be -1.
  """
  if isinstance(tensor, ops.EagerTensor):
    return tensor_shape.as_shape(
        [dim if dim != -1 else None for dim in tensor.numpy()])

  if tensor.get_shape().ndims == 0:
    value = constant_value(tensor)
    if value is None:
      raise ValueError(
          "Received a scalar with unknown value as shape; require a statically "
          "known scalar with value '-1' to describe an unknown shape.")
    if value != -1:
      raise ValueError(
          "Received a scalar value '%s' as shape; require a statically known "
          "scalar with value '-1' to describe an unknown shape." % value)
    return tensor_shape.unknown_shape()

  shape = tensor.get_shape().with_rank(1)
  if shape == [0]:
    return tensor_shape.scalar()
  elif tensor.op.type == "Shape":
    return tensor.op.inputs[0].get_shape()
  elif tensor.op.type == "Pack":
    ret = tensor_shape.scalar()  # Empty list.
    # Since we expect rank 1 inputs, Pack's axis must be zero, otherwise it
    # would not be rank 1.
    assert tensor.op.get_attr("axis") == 0
    for pack_input in tensor.op.inputs:
      # `pack_input` must be a scalar. Attempt to evaluate it, and append it
      # to `ret`.
      pack_input_val = constant_value(pack_input)
      if pack_input_val is None or pack_input_val < 0:
        new_dim = tensor_shape.Dimension(None)
      else:
        new_dim = tensor_shape.Dimension(pack_input_val)
      ret = ret.concatenate([new_dim])
    return ret
  elif tensor.op.type == "Concat":
    # We assume that `tensor.op.inputs[0]` evaluates to 0, as this is
    # the only legal value when concatenating vectors, and it will
    # have been checked by a previous shape function.
    ret = tensor_shape.scalar()  # Empty list.
    for concat_input in tensor.op.inputs[1:]:
      # `concat_input` must be a vector. Attempt to evaluate it as a shape,
      # and concatenate it with `ret`.
      ret = ret.concatenate(constant_value_as_shape(concat_input))
    return ret
  elif tensor.op.type == "ConcatV2":
    # We assume that `tensor.op.inputs[-1]` evaluates to 0, as this is
    # the only legal value when concatenating vectors, and it will
    # have been checked by a previous shape function.
    ret = tensor_shape.scalar()  # Empty list.
    for concat_input in tensor.op.inputs[:-1]:
      # `concat_input` must be a vector. Attempt to evaluate it as a shape,
      # and concatenate it with `ret`.
      ret = ret.concatenate(constant_value_as_shape(concat_input))
    return ret
  elif tensor.op.type == "StridedSlice":
    try:
      begin = constant_value(tensor.op.inputs[1])
      end = constant_value(tensor.op.inputs[2])
      strides = constant_value(tensor.op.inputs[3])
      if begin is not None and end is not None and strides is not None:
        begin = begin[0]
        end = end[0]
        strides = strides[0]
        begin_mask = tensor.op.get_attr("begin_mask")
        if begin_mask == 1:
          begin = None
        end_mask = tensor.op.get_attr("end_mask")
        if end_mask == 1:
          end = None

        ellipsis_mask = tensor.op.get_attr("ellipsis_mask")
        new_axis_mask = tensor.op.get_attr("new_axis_mask")
        shrink_axis_mask = tensor.op.get_attr("shrink_axis_mask")
        valid_attributes = (not ellipsis_mask and not new_axis_mask and
                            not shrink_axis_mask and (not begin_mask or
                                                      (begin_mask == 1)) and
                            (not end_mask or (end_mask == 1)))
        if valid_attributes:  # additional inputs not supported
          prev = constant_value_as_shape(tensor.op.inputs[0])
          prev = prev[begin:end:strides]
          ret = tensor_shape.TensorShape(prev)
          return ret

    except ValueError:  # Could come from get_attr or slicing prev.
      pass
    except TypeError:  # Could come from slicing prev.
      pass

  ret = tensor_shape.unknown_shape(shape.dims[0].value)
  value = constant_value(tensor)
  if value is not None:
    ret = ret.merge_with(
        tensor_shape.TensorShape([d if d >= 0 else None for d in value]))
  return ret
Ejemplo n.º 24
0
def map_fn(fn,
           elems,
           dtype=None,
           parallel_iterations=None,
           back_prop=True,
           swap_memory=False,
           infer_shape=True,
           name=None,
           fn_output_signature=None):
    """Transforms `elems` by applying `fn` to each element unstacked on axis 0.

  See also `tf.scan`.

  `map_fn` unstacks `elems` on axis 0 to obtain a sequence of elements;
  calls `fn` to transform each element; and then stacks the transformed
  values back together.

  #### Mapping functions with single-Tensor inputs and outputs

  If `elems` is a single tensor and `fn`'s signature is `tf.Tensor->tf.Tensor`,
  then `map_fn(fn, elems)` is equivalent to
  `tf.stack([fn(elem) for elem in tf.unstack(elems)])`.  E.g.:

  >>> tf.map_fn(fn=lambda t: tf.range(t, t + 3), elems=tf.constant([3, 5, 2]))
  <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
    array([[3, 4, 5],
           [5, 6, 7],
           [2, 3, 4]], dtype=int32)>

  `map_fn(fn, elems).shape = [elems.shape[0]] + fn(elems[0]).shape`.

  #### Mapping functions with multi-arity inputs and outputs

  `map_fn` also supports functions with multi-arity inputs and outputs:

  * If `elems` is a tuple (or nested structure) of tensors, then those tensors
    must all have the same outer-dimension size (`num_elems`); and `fn` is
    used to transform each tuple (or structure) of corresponding slices from
    `elems`.  E.g., if `elems` is a tuple `(t1, t2, t3)`, then `fn` is used to
    transform each tuple of slices `(t1[i], t2[i], t3[i])`
    (where `0 <= i < num_elems`).

  * If `fn` returns a tuple (or nested structure) of tensors, then the
    result is formed by stacking corresponding elements from those structures.

  #### Specifying `fn`'s output signature

  If `fn`'s input and output signatures are different, then the output
  signature must be specified using `fn_output_signature`.  (The input and
  output signatures are differ if their structures, dtypes, or tensor types do
  not match).  E.g.:

  >>> tf.map_fn(fn=tf.strings.length,  # input & output have different dtypes
  ...           elems=tf.constant(["hello", "moon"]),
  ...           fn_output_signature=tf.int32)
  <tf.Tensor: shape=(2,), dtype=int32, numpy=array([5, 4], dtype=int32)>
  >>> tf.map_fn(fn=tf.strings.join,  # input & output have different structures
  ...           elems=[tf.constant(['The', 'A']), tf.constant(['Dog', 'Cat'])],
  ...           fn_output_signature=tf.string)
  <tf.Tensor: shape=(2,), dtype=string,
   numpy=array([b'TheDog', b'ACat'], dtype=object)>

  `fn_output_signature` can be specified using any of the following:

  * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
  * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
  * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
  * A (possibly nested) tuple, list, or dict containing the above types.

  #### RaggedTensors

  `map_fn` supports `tf.RaggedTensor` inputs and outputs.  In particular:

  * If `elems` is a `RaggedTensor`, then `fn` will be called with each
    row of that ragged tensor.
    * If `elems` has only one ragged dimension, then the values passed to
      `fn` will be `tf.Tensor`s.
    * If `elems` has multiple ragged dimensions, then the values passed to
      `fn` will be `tf.RaggedTensor`s with one fewer ragged dimension.

  * If the result of `map_fn` should be a `RaggedTensor`, then use a
    `tf.RaggedTensorSpec` to specify `fn_output_signature`.
    * If `fn` returns `tf.Tensor`s with varying sizes, then use a
      `tf.RaggedTensorSpec` with `ragged_rank=0` to combine them into a
      single ragged tensor (which will have ragged_rank=1).
    * If `fn` returns `tf.RaggedTensor`s, then use a `tf.RaggedTensorSpec`
      with the same `ragged_rank`.

  >>> # Example: RaggedTensor input
  >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
  >>> tf.map_fn(tf.reduce_sum, rt, fn_output_signature=tf.int32)
  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([6, 0, 9, 6], dtype=int32)>

  >>> # Example: RaggedTensor output
  >>> elems = tf.constant([3, 5, 0, 2])
  >>> tf.map_fn(tf.range, elems,
  ...           fn_output_signature=tf.RaggedTensorSpec(shape=[None],
  ...                                                   dtype=tf.int32))
  <tf.RaggedTensor [[0, 1, 2], [0, 1, 2, 3, 4], [], [0, 1]]>

  Note: `map_fn` should only be used if you need to map a function over the
  *rows* of a `RaggedTensor`.  If you wish to map a function over the
  individual values, then you should use:

  * `tf.ragged.map_flat_values(fn, rt)`
    (if fn is expressible as TensorFlow ops)
  * `rt.with_flat_values(map_fn(fn, rt.flat_values))`
    (otherwise)

  E.g.:

  >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
  >>> tf.ragged.map_flat_values(lambda x: x + 2, rt)
  <tf.RaggedTensor [[3, 4, 5], [], [6, 7], [8]]>

  #### SparseTensors

  `map_fn` supports `tf.sparse.SparseTensor` inputs and outputs.  In particular:

  * If `elems` is a `SparseTensor`, then `fn` will be called with each row
    of that sparse tensor. In particular, the value passed to `fn` will be a
    `tf.sparse.SparseTensor` with one fewer dimension than `elems`.

  * If the result of `map_fn` should be a `SparseTensor`, then use a
    `tf.SparseTensorSpec` to specify `fn_output_signature`.  The individual
    `SparseTensor`s returned by `fn` will be stacked into a single
    `SparseTensor` with one more dimension.

  >>> # Example: SparseTensor input
  >>> st = tf.sparse.SparseTensor([[0, 0], [2, 0], [2, 1]], [2, 3, 4], [4, 4])
  >>> tf.map_fn(tf.sparse.reduce_sum, st, fn_output_signature=tf.int32)
  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([2, 0, 7, 0], dtype=int32)>

  >>> # Example: SparseTensor output
  >>> tf.sparse.to_dense(
  ...     tf.map_fn(tf.sparse.eye, tf.constant([2, 3]),
  ...               fn_output_signature=tf.SparseTensorSpec(None, tf.float32)))
  <tf.Tensor: shape=(2, 3, 3), dtype=float32, numpy=
    array([[[1., 0., 0.],
            [0., 1., 0.],
            [0., 0., 0.]],
           [[1., 0., 0.],
            [0., 1., 0.],
            [0., 0., 1.]]], dtype=float32)>

  Note: `map_fn` should only be used if you need to map a function over the
  *rows* of a `SparseTensor`.  If you wish to map a function over the nonzero
  values, then you should use:

  * If the function is expressible as TensorFlow ops, use:
    ```python
    tf.sparse.SparseTensor(st.indices, fn(st.values), st.dense_shape)
    ```
  * Otherwise, use:
    ```python
    tf.sparse.SparseTensor(st.indices, tf.map_fn(fn, st.values),
                           st.dense_shape)
    ```

  #### `map_fn` vs. vectorized operations

  `map_fn` will apply the operations used by `fn` to each element of `elems`,
  resulting in `O(elems.shape[0])` total operations.  This is somewhat
  mitigated by the fact that `map_fn` can process elements in parallel.
  However, a transform expressed using `map_fn` is still typically less
  efficient than an equivalent transform expressed using vectorized operations.

  `map_fn` should typically only be used if one of the following is true:

  * It is difficult or expensive to express the desired transform with
    vectorized operations.
  * `fn` creates large intermediate values, so an equivalent vectorized
    transform would take too much memory.
  * Processing elements in parallel is more efficient than an equivalent
    vectorized transform.
  * Efficiency of the transform is not critical, and using `map_fn` is
    more readable.

  E.g., the example given above that maps `fn=lambda t: tf.range(t, t + 3)`
  across `elems` could be rewritten more efficiently using vectorized ops:

  >>> elems = tf.constant([3, 5, 2])
  >>> tf.range(3) + tf.expand_dims(elems, 1)
  <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
    array([[3, 4, 5],
           [5, 6, 7],
           [2, 3, 4]], dtype=int32)>

  In some cases, `tf.vectorized_map` can be used to automatically convert a
  function to a vectorized equivalent.

  #### Eager execution

  When executing eagerly, `map_fn` does not execute in parallel even if
  `parallel_iterations` is set to a value > 1. You can still get the
  performance benefits of running a function in parallel by using the
  `tf.function` decorator:

  >>> fn=lambda t: tf.range(t, t + 3)
  >>> @tf.function
  ... def func(elems):
  ...   return tf.map_fn(fn, elems, parallel_iterations=3)
  >>> func(tf.constant([3, 5, 2]))
  <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
    array([[3, 4, 5],
           [5, 6, 7],
           [2, 3, 4]], dtype=int32)>


  Note: if you use the `tf.function` decorator, any non-TensorFlow Python
  code that you may have written in your function won't get executed. See
  `tf.function` for more  details. The recommendation would be to debug without
  `tf.function` but switch to it to get performance benefits of running `map_fn`
  in parallel.

  Args:
    fn: The callable to be performed.  It accepts one argument, which will have
      the same (possibly nested) structure as `elems`.  Its output must have the
      same structure as `fn_output_signature` if one is provided; otherwise it
      must have the same structure as `elems`.
    elems: A tensor or (possibly nested) sequence of tensors, each of which will
      be unstacked along their first dimension.  `fn` will be applied to the
      nested sequence of the resulting slices.  `elems` may include ragged and
      sparse tensors. `elems` must consist of at least one tensor.
    dtype: Deprecated: Equivalent to `fn_output_signature`.
    parallel_iterations: (optional) The number of iterations allowed to run in
      parallel. When graph building, the default value is 10. While executing
      eagerly, the default value is set to 1.
    back_prop: (optional) False disables support for back propagation.
    swap_memory: (optional) True enables GPU-CPU memory swapping.
    infer_shape: (optional) False disables tests for consistent output shapes.
    name: (optional) Name prefix for the returned tensors.
    fn_output_signature: The output signature of `fn`. Must be specified if
      `fn`'s input and output signatures are different (i.e., if their
      structures, dtypes, or tensor types do not match).
      `fn_output_signature` can be specified using any of the following:

      * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
      * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
      * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
      * A (possibly nested) tuple, list, or dict containing the above types.

  Returns:
    A tensor or (possibly nested) sequence of tensors.  Each tensor stacks the
    results of applying `fn` to tensors unstacked from `elems` along the first
    dimension, from first to last.  The result may include ragged and sparse
    tensors.

  Raises:
    TypeError: if `fn` is not callable or the structure of the output of
      `fn` and `fn_output_signature` do not match.
    ValueError: if the lengths of the output of `fn` and `fn_output_signature`
      do not match, or if the `elems` does not contain any tensor.

  Examples:

    >>> elems = np.array([1, 2, 3, 4, 5, 6])
    >>> tf.map_fn(lambda x: x * x, elems)
    <tf.Tensor: shape=(6,), dtype=int64, numpy=array([ 1,  4,  9, 16, 25, 36])>

    >>> elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
    >>> tf.map_fn(lambda x: x[0] * x[1], elems, fn_output_signature=tf.int64)
    <tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1,  2, -3])>

    >>> elems = np.array([1, 2, 3])
    >>> tf.map_fn(lambda x: (x, -x), elems,
    ...          fn_output_signature=(tf.int64, tf.int64))
    (<tf.Tensor: shape=(3,), dtype=int64, numpy=array([1, 2, 3])>,
     <tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1, -2, -3])>)
  """
    # This function uses a `while_loop` to call `fn` on each value of the input
    # tensor(s) (unstacked on dimension 0).  The following sequence of variables
    # are used to transform the input tensor(s) (`elems`) into the output
    # tensor(s) (`result`):
    #
    #   - Preparing and unstacking input values for the while_loop:
    #     - elems: The input tensor(s) to map_fn. May include composite tensors.
    #     - elems_flat: Flattened list of tensors from elems (using nest.flatten)
    #                   May include composite tensors.
    #     - elems_batchable: Concatenation of "batchable tensor lists" for each
    #                        tensor in elems_flat.  This "boxes" composite tensors
    #                        into sliceable tf.Tensor objects.  For more info see:
    #                        TensorSpec._to_batched_tensor_list
    #     - elems_batchable_ta: List of TensorArrays used to unstack each Tensor
    #                           in elems_batchable into elems_value_batchable.
    #
    #   - Calling `fn` on each unstacked value in the body of the while_loop:
    #     - elems_value_batchable: Single unstacked value from elems_batchable.
    #     - elems_value_flat: Single unstacked value from elems_flat,
    #                         constructed from elems_value_batchable (using
    #                         TensorSpec._from_tensor_list).
    #     - elems_value: Single unstacked value from elems (the input to fn).
    #     - result_value: Result of calling `fn(elems_value)`.  May contain
    #                     composite tensors.
    #     - result_value_flat: Flattened list of tensors from result_value.
    #                          May contain composite tensors.
    #     - result_value_batchable: Concatenation of batchable tensor lists for
    #                               each tensor in result_value_flat
    #                               (using TensorSpec._to_tensor_list).
    #
    #   - Collecting and stacking output values from the while_loop:
    #     - result_batchable_ta: List of TensorArrays used to stack each tensor
    #                            ta result_value_batchable into result_batchable.
    #     - result_batchable: Stacked tensors from result_batchable_ta.
    #     - result_flat: Flat list of tensors for the result, constructed from
    #                    results bactchable (using TensorSpec._from_tensor_list).
    #     - result: Structured result value packed from results flat
    #               (using nest.pack_sequence_as).

    if fn_output_signature is None:
        fn_output_signature = dtype

    if not callable(fn):
        raise TypeError(f"The provided function {fn.__name__} is not callable."
                        "fn must be callable.")

    in_graph_mode = not context.executing_eagerly()
    # Set the default number of parallel_iterations depending on graph/eager mode.
    if in_graph_mode and not parallel_iterations:
        parallel_iterations = 10
    elif not in_graph_mode and not parallel_iterations:
        parallel_iterations = 1
    elif not in_graph_mode and parallel_iterations > 1:
        logging.log_first_n(
            logging.WARN, "Setting parallel_iterations > 1 has no "
            "effect when executing eagerly. Consider calling map_fn"
            " with tf.function to execute fn in "
            "parallel.", 1)
        parallel_iterations = 1

    # Flatten the input tensors, and get the TypeSpec for each one.
    elems_flat = nest.flatten(elems)

    # Check in case this is an empty list
    if len(elems_flat) == 0:
        raise ValueError(
            "elems must be a Tensor or (possibly nested) sequence of Tensors. "
            "Got {}, which does not contain any Tensors.".format(elems))

    elems_flat_signature = [
        type_spec.type_spec_from_value(e) for e in elems_flat
    ]
    elems_unflatten = lambda x: nest.pack_sequence_as(elems, x)

    # Flatten fn's output signature.
    if fn_output_signature is None:
        # If fn_output_signature was not specified, then assume that it matches the
        # input signature.
        result_flat_signature = [
            _most_general_compatible_type(s)._unbatch()  # pylint: disable=protected-access
            for s in elems_flat_signature
        ]
        result_unflatten = elems_unflatten
    else:
        result_flat_signature = [
            _dtype_to_spec(d) for d in nest.flatten(fn_output_signature)
        ]
        result_unflatten = lambda x: nest.pack_sequence_as(
            fn_output_signature, x)

    with ops.name_scope(name, "map", elems_flat):
        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode:
            # Any get_variable calls in fn will cache the first call locally
            # and not issue repeated network I/O requests for each iteration.
            varscope = vs.get_variable_scope()
            varscope_caching_device_was_none = False
            if varscope.caching_device is None:
                # TODO(ebrevdo): Change to using colocate_with here and in other
                # methods.
                varscope.set_caching_device(lambda op: op.device)
                varscope_caching_device_was_none = True

        elems_flat = [
            ops.convert_to_tensor_or_composite(t, name="elem")
            for t in elems_flat
        ]

        # Check that inputs are not scalars.
        first_elem = elems_flat[0]
        if hasattr(first_elem, "shape"):
            elems_static_shape = first_elem.shape
            if elems_static_shape.ndims is not None and elems_static_shape.ndims < 1:
                raise ValueError(
                    "Elements in elems must be 1+ dimensional Tensors, not scalars"
                )

        # Box any composite tensors into tensor lists.
        elems_batchable = _elems_flat_to_batchable(elems_flat)

        # Find the number of iterations, n.  (may be known statically.)
        n_static = tensor_shape.Dimension(
            tensor_shape.dimension_value(
                elems_batchable[0].get_shape().with_rank_at_least(1)[0]))
        for tensor in elems_batchable[1:]:
            n_static.assert_is_compatible_with(
                tensor_shape.Dimension(
                    tensor_shape.dimension_value(
                        tensor.get_shape().with_rank_at_least(1)[0])))
        n = n_static.value or array_ops.shape(elems_batchable[0])[0]

        # Convert elems to tensor array.
        # TODO(edloper): Should we set infer_shape=False for composite tensors?
        elems_batchable_ta = [
            tensor_array_ops.TensorArray(dtype=t.dtype,
                                         size=n,
                                         dynamic_size=False,
                                         infer_shape=True)
            for t in elems_batchable
        ]
        # Unpack elements
        elems_batchable_ta = [
            ta.unstack(t)
            for (ta, t) in zip(elems_batchable_ta, elems_batchable)
        ]

        i = constant_op.constant(0)

        # Prepare result tensor array.
        # TODO(edloper): Should we set infer_shape=False for composite tensors?
        result_batchable_tensor_spec = (
            _result_flat_signature_to_batchable_tensor_spec(
                result_flat_signature))
        result_batchable_ta = []
        for spec in result_batchable_tensor_spec:
            result_batchable_ta.append(
                tensor_array_ops.TensorArray(dtype=spec.dtype,
                                             size=n,
                                             dynamic_size=False,
                                             infer_shape=infer_shape,
                                             element_shape=spec.shape))

        def compute(i, tas):
            """The loop body of map_fn.

      Args:
        i: the loop counter
        tas: the flat TensorArray accumulator list

      Returns:
        (i + 1, tas): the updated counter + updated TensorArrays

      Raises:
        TypeError: if fn_output_signature and result_value structure don't match
        ValueType: if fn_output_signature and result_value lengths don't match
      """
            elems_value_batchable = [ta.read(i) for ta in elems_batchable_ta]
            elems_value_flat = _elems_value_batchable_to_flat(
                elems_value_batchable, elems_flat_signature)
            elems_value = elems_unflatten(elems_value_flat)
            ag_ctx = autograph_ctx.control_status_ctx()
            autographed_fn = autograph.tf_convert(fn, ag_ctx)
            result_value = autographed_fn(elems_value)
            nest.assert_same_structure(fn_output_signature or elems,
                                       result_value)
            result_value_flat = nest.flatten(result_value)
            result_value_batchable = _result_value_flat_to_batchable(
                result_value_flat, result_flat_signature)
            tas = [
                ta.write(i, value)
                for (ta, value) in zip(tas, result_value_batchable)
            ]
            return (i + 1, tas)

        _, r_a = control_flow_ops.while_loop(
            lambda i, _: i < n,
            compute, (i, result_batchable_ta),
            parallel_iterations=parallel_iterations,
            back_prop=back_prop,
            swap_memory=swap_memory,
            maximum_iterations=n)
        result_batchable = [r.stack() for r in r_a]

        # Update each output tensor w/ static shape info about the outer dimension.
        for r in result_batchable:
            r.set_shape(
                tensor_shape.TensorShape(n_static).concatenate(
                    r.get_shape()[1:]))

        # TODO(akshayka): Remove the in_graph_mode check once caching devices are
        # supported in Eager
        if in_graph_mode and varscope_caching_device_was_none:
            varscope.set_caching_device(None)

        result_flat = _result_batchable_to_flat(result_batchable,
                                                result_flat_signature,
                                                n_static)
        result = result_unflatten(result_flat)
        return result
Ejemplo n.º 25
0
 def testTruedivFails(self):
     unknown = tensor_shape.Dimension(None)
     self.assertEqual((unknown // unknown).value, None)
     with self.assertRaisesRegexp(TypeError, r"unsupported operand type"):
         unknown / unknown  # pylint: disable=pointless-statement
Ejemplo n.º 26
0
def constant_value_as_shape(tensor):  # pylint: disable=invalid-name
    """A version of `constant_value()` that returns a `TensorShape`.

  This version should be used when a constant tensor value is
  interpreted as a (possibly partial) shape, e.g. in the shape
  function for `tf.reshape()`. By explicitly requesting a
  `TensorShape` as the return value, it is possible to represent
  unknown dimensions; by contrast, `constant_value()` is
  all-or-nothing.

  Args:
    tensor: The rank-0 or rank-1 Tensor to be evaluated.

  Returns:
    A `TensorShape` based on the constant value of the given `tensor`.

  Raises:
    ValueError: If the shape is rank-0 and is not statically known to be -1.
  """
    if isinstance(tensor, ops.EagerTensor):
        return tensor_shape.as_shape(
            [dim if dim != -1 else None for dim in tensor.numpy()])

    if tensor.get_shape().ndims == 0:
        value = constant_value(tensor)
        if value is None:
            raise ValueError(
                "Received a scalar with unknown value as shape; require a statically "
                "known scalar with value '-1' to describe an unknown shape.")
        if value != -1:
            raise ValueError(
                "Received a scalar value '%s' as shape; require a statically known "
                "scalar with value '-1' to describe an unknown shape." % value)
        return tensor_shape.unknown_shape()

    shape = tensor.get_shape().with_rank(1)
    if shape == [0]:
        return tensor_shape.TensorShape([])
    elif tensor.op.type == "Cast":
        pre_cast = constant_value_as_shape(tensor.op.inputs[0])
        if pre_cast.dims is None:
            # the input to cast has a totally undefined shape; just return that.
            return pre_cast
        cast_dtype = dtypes.as_dtype(tensor.op.get_attr("DstT"))
        if cast_dtype not in (dtypes.int32, dtypes.int64):
            return tensor_shape.unknown_shape(shape.dims[0].value)
        dest_dtype_shape_array = np.array([
            x if x is not None else -1 for x in pre_cast.as_list()
        ]).astype(cast_dtype.as_numpy_dtype)
        return tensor_shape.TensorShape(
            [x if x >= 0 else None for x in dest_dtype_shape_array])
    elif tensor.op.type == "Shape":
        return tensor.op.inputs[0].get_shape()
    elif tensor.op.type == "Pack":
        ret = tensor_shape.TensorShape([])  # Empty list.
        # Since we expect rank 1 inputs, Pack's axis must be zero, otherwise it
        # would not be rank 1.
        assert tensor.op.get_attr("axis") == 0
        for pack_input in tensor.op.inputs:
            # `pack_input` must be a scalar. Attempt to evaluate it, and append it
            # to `ret`.
            pack_input_val = constant_value(pack_input)
            if pack_input_val is None or pack_input_val < 0:
                new_dim = tensor_shape.Dimension(None)
            else:
                new_dim = tensor_shape.Dimension(pack_input_val)
            ret = ret.concatenate([new_dim])
        return ret
    elif tensor.op.type == "Concat":
        # We assume that `tensor.op.inputs[0]` evaluates to 0, as this is
        # the only legal value when concatenating vectors, and it will
        # have been checked by a previous shape function.
        ret = tensor_shape.TensorShape([])  # Empty list.
        for concat_input in tensor.op.inputs[1:]:
            # `concat_input` must be a vector. Attempt to evaluate it as a shape,
            # and concatenate it with `ret`.
            ret = ret.concatenate(constant_value_as_shape(concat_input))
        return ret
    elif tensor.op.type == "ConcatV2":
        # We assume that `tensor.op.inputs[-1]` evaluates to 0, as this is
        # the only legal value when concatenating vectors, and it will
        # have been checked by a previous shape function.
        ret = tensor_shape.TensorShape([])  # Empty list.
        for concat_input in tensor.op.inputs[:-1]:
            # `concat_input` must be a vector. Attempt to evaluate it as a shape,
            # and concatenate it with `ret`.
            ret = ret.concatenate(constant_value_as_shape(concat_input))
        return ret
    elif tensor.op.type == "StridedSlice":
        try:
            begin = constant_value(tensor.op.inputs[1])
            end = constant_value(tensor.op.inputs[2])
            strides = constant_value(tensor.op.inputs[3])
            if begin is not None and end is not None and strides is not None:
                begin = begin[0]
                end = end[0]
                strides = strides[0]
                begin_mask = tensor.op.get_attr("begin_mask")
                if begin_mask == 1:
                    begin = None
                end_mask = tensor.op.get_attr("end_mask")
                if end_mask == 1:
                    end = None

                ellipsis_mask = tensor.op.get_attr("ellipsis_mask")
                new_axis_mask = tensor.op.get_attr("new_axis_mask")
                shrink_axis_mask = tensor.op.get_attr("shrink_axis_mask")
                valid_attributes = (not ellipsis_mask and not new_axis_mask
                                    and not shrink_axis_mask
                                    and (not begin_mask or (begin_mask == 1))
                                    and (not end_mask or (end_mask == 1)))
                if valid_attributes:  # additional inputs not supported
                    prev = constant_value_as_shape(tensor.op.inputs[0])
                    prev = prev[begin:end:strides]
                    ret = tensor_shape.TensorShape(prev)
                    return ret

        except ValueError:  # Could come from get_attr or slicing prev.
            pass
        except TypeError:  # Could come from slicing prev.
            pass
    elif (tensor.op.type == "Placeholder" and tensor.op.graph.building_function
          and hasattr(tensor.op.graph, "internal_captures")):
        # If we are inside a FuncGraph try to lookup the constant value of the
        # corresponding external capture. Note that we only look at captures and
        # not the fed inputs because those can be fed different values in different
        # instantiations of the function call or different iterations of a
        # tf.while_loop.
        for i, capture in enumerate(tensor.op.graph.internal_captures):
            if capture is tensor:
                external_capture = tensor.op.graph.external_captures[i]
                return constant_value_as_shape(external_capture)

    ret = tensor_shape.unknown_shape(shape.dims[0].value)
    value = constant_value(tensor)
    if value is not None:
        ret = ret.merge_with(
            tensor_shape.TensorShape([d if d >= 0 else None for d in value]))
    return ret
Ejemplo n.º 27
0
 def testKnownAndUnknownDimensions(self):
     known = tensor_shape.Dimension(12)
     unknown = tensor_shape.Dimension(None)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (known + unknown).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (unknown + known).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (known * unknown).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (unknown * known).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (known // unknown).value)
     self.assertEqual(
         tensor_shape.Dimension(None).value, (unknown // known).value)
     self.assertEqual(tensor_shape.Dimension(12), known.merge_with(unknown))
     self.assertEqual(tensor_shape.Dimension(12), unknown.merge_with(known))
     self.assertIs(
         None,
         tensor_shape.Dimension(12) < tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(12) <= tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(12) > tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(12) >= tensor_shape.Dimension(None))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) < tensor_shape.Dimension(12))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) <= tensor_shape.Dimension(12))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) > tensor_shape.Dimension(12))
     self.assertIs(
         None,
         tensor_shape.Dimension(None) >= tensor_shape.Dimension(12))
Ejemplo n.º 28
0
  def testInequality(self):
    self.assertTrue(tensor_shape.Dimension(12) != tensor_shape.Dimension(13))
    self.assertFalse(tensor_shape.Dimension(12) != tensor_shape.Dimension(12))
    self.assertIs(None,
                  tensor_shape.Dimension(12) != tensor_shape.Dimension(None))
    self.assertIs(None,
                  tensor_shape.Dimension(None) != tensor_shape.Dimension(12))
    self.assertIs(None,
                  tensor_shape.Dimension(None) != tensor_shape.Dimension(None))

    # None indicates ambiguous comparison, but comparison vs the wrong type
    # is unambigously False.
    self.assertIsNotNone(tensor_shape.Dimension(12) != "_")
    self.assertIsNotNone(tensor_shape.Dimension(None) != 12.99)
    self.assertTrue(tensor_shape.Dimension(12) != "_")
    self.assertTrue(tensor_shape.Dimension(None) != 12.99)

    self.assertIs(None, tensor_shape.Dimension(None) != "13")
    self.assertIs(None, tensor_shape.Dimension(None) != None)  # pylint: disable=g-equals-none
    self.assertTrue(tensor_shape.Dimension(12) != 12.99)
Ejemplo n.º 29
0
    def batch_jacobian(self,
                       target,
                       source,
                       unconnected_gradients=UnconnectedGradients.NONE,
                       parallel_iterations=None,
                       experimental_use_pfor=True):
        """Computes and stacks per-example jacobians.

    See http://en.wikipedia.org/wiki/jacobian_matrix_and_determinant for the
    definition of a Jacobian.  This function is essentially an efficient
    implementation of the following:
    `tf.stack([self.jacobian(y[i], x[i]) for i in range(x.shape[0])])`.

    Note that compared to `GradientTape.jacobian` which computes gradient of
    each output value w.r.t each input value, this function is useful when
    `target[i,...] is independent of `source[j,...]` for `j != i`. This
    independence assumption allows more efficient computation as compared to
    `GradientTape.jacobian`. The output, as well as intermediate activations,
    are lower dimensional and avoid a bunch of redundant zeros which would
    result in the jacobian computation given the independence assumption.

    Example usage:
    ```python
    with tf.GradientTape() as g:
      x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
      g.watch(x)
      y = x * x
    batch_jacobian = g.batch_jacobian(y, x)
    # batch_jacobian is [[[2,  0], [0,  4]], [[6,  0], [0,  8]]]
    ```

    Args:
      target: A tensor with rank 2 or higher and with shape [b, y1, ..., y_n].
        `target[i,...]` should only depend on `source[i,...]`.
      source: A tensor with rank 2 or higher and with shape [b, x1, ..., x_m].
      unconnected_gradients: a value which can either hold 'none' or 'zero' and
        alters the value which will be returned if the target and sources are
        unconnected. The possible values and effects are detailed in
        'UnconnectedGradients' and it defaults to 'none'.
      parallel_iterations: A knob to control how many iterations are dispatched
        in parallel. This knob can be used to control the total memory usage.
      experimental_use_pfor: If true, uses pfor for computing the Jacobian. Else
        uses a tf.while_loop.

    Returns:
      A tensor `t` with shape [b, y_1, ..., y_n, x1, ..., x_m] where `t[i, ...]`
      is the jacobian of `target[i, ...]` w.r.t. `source[i, ...]`, i.e. stacked
      per-example jacobians.

    Raises:
      RuntimeError: If called on a non-persistent tape with eager execution
        enabled and without enabling experimental_use_pfor.
      ValueError: If vectorization of jacobian computation fails or if first
        dimension of `target` and `source` do not match.
    """
        target_shape = target.shape
        if target_shape.rank is None:
            dim = tensor_shape.Dimension(None)
        else:
            dim = target_shape.dims[0]
        if not (target_shape.with_rank_at_least(2)
                and source.shape.with_rank_at_least(2)
                and dim.is_compatible_with(source.shape[0])):
            raise ValueError("Need first dimension of target shape (%s) and "
                             "source shape (%s) to match." %
                             (target.shape, source.shape))
        if target_shape.is_fully_defined():
            batch_size = int(target_shape[0])
            target_row_size = target_shape.num_elements() // batch_size
        else:
            target_shape = array_ops.shape(target)
            batch_size = target_shape[0]
            target_row_size = array_ops.size(target) // batch_size
        source_shape = array_ops.shape(source)
        # Flatten target to 2-D.
        # Note that we push and pop the tape here and below. This is needed since we
        # need gradients through the enclosed operations.
        self._push_tape()
        with ops.control_dependencies(
            [check_ops.assert_equal(batch_size, source_shape[0])]):
            target = array_ops.reshape(target, [batch_size, target_row_size])
        self._pop_tape()

        def loop_fn(i):
            self._push_tape()
            y = array_ops.gather(target, i, axis=1)
            self._pop_tape()
            return self.gradient(y,
                                 source,
                                 unconnected_gradients=unconnected_gradients)

        if experimental_use_pfor:
            try:
                output = pfor_ops.pfor(loop_fn,
                                       target_row_size,
                                       parallel_iterations=parallel_iterations)
            except ValueError as err:
                six.reraise(
                    ValueError,
                    ValueError(
                        str(err) +
                        "\nEncountered an exception while vectorizing the "
                        "batch_jacobian computation. Vectorization can be disabled by "
                        "setting experimental_use_pfor to False."),
                    sys.exc_info()[2])
        else:
            if context.executing_eagerly() and not self._persistent:
                raise RuntimeError(
                    "GradientTape must be created with persistent=True"
                    " to compute the batch_jacobian with eager execution enabled and "
                    " with experimental_use_pfor set to False.")
            output = pfor_ops.for_loop(loop_fn,
                                       target.dtype,
                                       target_row_size,
                                       parallel_iterations=parallel_iterations)
        if output is None:
            return None
        output = array_ops.reshape(output, [target_row_size, batch_size, -1])
        output = array_ops.transpose(output, [1, 0, 2])
        new_shape = array_ops.concat([target_shape, source_shape[1:]], axis=0)
        return array_ops.reshape(output, new_shape)
Ejemplo n.º 30
0
 def testRepr(self):
   self.assertEqual(repr(tensor_shape.Dimension(7)), "Dimension(7)")
   self.assertEqual(repr(tensor_shape.Dimension(None)), "Dimension(None)")