Ejemplo n.º 1
0
def clip_by_value(t, clip_value_min, clip_value_max,
                  name=None):
  """Clips tensor values to a specified min and max.

  Given a tensor `t`, this operation returns a tensor of the same type and
  shape as `t` with its values clipped to `clip_value_min` and `clip_value_max`.
  Any values less than `clip_value_min` are set to `clip_value_min`. Any values
  greater than `clip_value_max` are set to `clip_value_max`.

  Args:
    t: A `Tensor`.
    clip_value_min: A 0-D (scalar) `Tensor`. The minimum value to clip by.
    clip_value_max: A 0-D (scalar) `Tensor`. The maximum value to clip by.
    name: A name for the operation (optional).

  Returns:
    A clipped `Tensor`.
  """
  with ops.op_scope([t, clip_value_min, clip_value_max], name,
                   "clip_by_value") as name:
    t = ops.convert_to_tensor(t, name="t")

    # Go through list of tensors, for each value in each tensor clip
    t_min = math_ops.minimum(
        t, array_ops.fill(array_ops.shape(t), clip_value_max))
    t_max = math_ops.maximum(
        t_min, array_ops.fill(array_ops.shape(t), clip_value_min),
        name=name)

  return t_max
  def testParallelUpdateWithLocking(self):
    with self.test_session() as sess:
      zeros_t = array_ops.fill([1024, 1024], 0.0)
      ones_t = array_ops.fill([1024, 1024], 1.0)
      p = variables.Variable(zeros_t)
      adds = [
          state_ops.assign_add(
              p, ones_t, use_locking=True) for _ in range(20)
      ]
      p.initializer.run()

      def run_add(add_op):
        sess.run(add_op)

      threads = [
          self.checkedThread(
              target=run_add, args=(add_op,)) for add_op in adds
      ]
      for t in threads:
        t.start()
      for t in threads:
        t.join()

      vals = p.eval()
      ones = np.ones((1024, 1024)).astype(np.float32)
      self.assertAllEqual(vals, ones * 20)
Ejemplo n.º 3
0
  def _variance(self):
    # We need to put the tf.where inside the outer tf.where to ensure we never
    # hit a NaN in the gradient.
    denom = array_ops.where(math_ops.greater(self.df, 2.),
                            self.df - 2.,
                            array_ops.ones_like(self.df))
    # Abs(scale) superfluous.
    var = (array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype) *
           math_ops.square(self.scale) * self.df / denom)
    # When 1 < df <= 2, variance is infinite.
    inf = np.array(np.inf, dtype=self.dtype.as_numpy_dtype())
    result_where_defined = array_ops.where(
        self.df > array_ops.fill(self.batch_shape_tensor(), 2.),
        var,
        array_ops.fill(self.batch_shape_tensor(), inf, name="inf"))

    if self.allow_nan_stats:
      nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
      return array_ops.where(
          math_ops.greater(
              self.df,
              array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)),
          result_where_defined,
          array_ops.fill(self.batch_shape_tensor(), nan, name="nan"))
    else:
      return control_flow_ops.with_dependencies(
          [
              check_ops.assert_less(
                  array_ops.ones([], dtype=self.dtype),
                  self.df,
                  message="variance not defined for components of df <= 1"),
          ],
          result_where_defined)
  def testParallelUpdateWithLocking(self):
    # We need each thread to keep its own device stack or the device scopes
    # won't be properly nested.
    ops.get_default_graph().switch_to_thread_local()
    with self.cached_session() as sess:
      zeros_t = array_ops.fill([1024, 1024], 0.0)
      ones_t = array_ops.fill([1024, 1024], 1.0)
      p = variables.Variable(zeros_t)
      adds = [
          state_ops.assign_add(
              p, ones_t, use_locking=True) for _ in range(20)
      ]
      self.evaluate(p.initializer)

      def run_add(add_op):
        self.evaluate(add_op)

      threads = [
          self.checkedThread(
              target=run_add, args=(add_op,)) for add_op in adds
      ]
      for t in threads:
        t.start()
      for t in threads:
        t.join()

      vals = self.evaluate(p)
      ones = np.ones((1024, 1024)).astype(np.float32)
      self.assertAllEqual(vals, ones * 20)
  def testParallelAssignWithLocking(self):
    with self.test_session() as sess:
      zeros_t = array_ops.fill([1024, 1024], 0.0)
      ones_t = array_ops.fill([1024, 1024], 1.0)
      p = variables.Variable(zeros_t)
      assigns = [
          state_ops.assign(
              p, math_ops.mul(ones_t, float(i)), use_locking=True)
          for i in range(1, 21)
      ]
      p.initializer.run()

      def run_assign(assign_op):
        sess.run(assign_op)

      threads = [
          self.checkedThread(
              target=run_assign, args=(assign_op,)) for assign_op in assigns
      ]
      for t in threads:
        t.start()
      for t in threads:
        t.join()

      vals = p.eval()

      # Assert every element is the same, and taken from one of the assignments.
      self.assertTrue(vals[0, 0] > 0)
      self.assertTrue(vals[0, 0] <= 20)
      self.assertAllEqual(vals, np.ones([1024, 1024]) * vals[0, 0])
  def testParallelAssignWithLocking(self):
    # We need each thread to keep its own device stack or the device scopes
    # won't be properly nested.
    ops.get_default_graph().switch_to_thread_local()
    with self.cached_session() as sess:
      zeros_t = array_ops.fill([1024, 1024], 0.0)
      ones_t = array_ops.fill([1024, 1024], 1.0)
      p = variables.Variable(zeros_t)
      assigns = [
          state_ops.assign(
              p, math_ops.multiply(ones_t, float(i)), use_locking=True)
          for i in range(1, 21)
      ]
      self.evaluate(p.initializer)

      def run_assign(assign_op):
        self.evaluate(assign_op)

      threads = [
          self.checkedThread(
              target=run_assign, args=(assign_op,)) for assign_op in assigns
      ]
      for t in threads:
        t.start()
      for t in threads:
        t.join()

      vals = self.evaluate(p)

      # Assert every element is the same, and taken from one of the assignments.
      self.assertTrue(vals[0, 0] > 0)
      self.assertTrue(vals[0, 0] <= 20)
      self.assertAllEqual(vals, np.ones([1024, 1024]) * vals[0, 0])
Ejemplo n.º 7
0
    def _variance(self):
        var = self._ones() * math_ops.square(self.sigma) * self.df / (self.df - 2)
        # When 1 < df <= 2, variance is infinite.
        inf = np.array(np.inf, dtype=self.dtype.as_numpy_dtype())
        result_where_defined = math_ops.select(
            math_ops.greater(self.df, array_ops.fill(self.batch_shape(), 2.0)),
            var,
            array_ops.fill(self.batch_shape(), inf, name="inf"),
        )

        if self.allow_nan_stats:
            nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
            return math_ops.select(
                math_ops.greater(self.df, self._ones()),
                result_where_defined,
                array_ops.fill(self.batch_shape(), nan, name="nan"),
            )
        else:
            return control_flow_ops.with_dependencies(
                [
                    check_ops.assert_less(
                        array_ops.ones((), dtype=self.dtype),
                        self.df,
                        message="variance not defined for components of df <= 1",
                    )
                ],
                result_where_defined,
            )
Ejemplo n.º 8
0
def _ConcatGrad(op, grad):
  """Gradient for concat op."""
  assert isinstance(grad, ops.Tensor)
  # Degenerate concatenation, just return grad.
  if len(op.inputs) == 2:
    return [None, grad]
  # Get the inputs' tensor shapes
  sizes = [array_ops.shape(x) for x in op.inputs[1:]]
  concat_dim = op.inputs[0]
  # Since shape is 1-D, shape_of_shape = [rank-of-inputs]
  shape_of_shape = array_ops.shape(sizes[0])
  # Make a vector of length equal to the input's dimensions,
  # with 0's everywhere and 1 in the concat dim position.
  # Note: Can't use sparse_to_dense since it isn't GPU-capable (for now)
  mask = array_ops.concat(0,
                          [array_ops.fill(
                              array_ops.expand_dims(concat_dim, 0), 0), [1],
                           array_ops.fill(shape_of_shape - concat_dim - 1, 0)])
  out_grads = []
  begin = array_ops.fill(shape_of_shape, 0)
  for i in range(len(sizes)):
    out_grads.append(array_ops.slice(grad, begin, sizes[i]))
    # Lint complains begin = begin + ...
    begin = math_ops.add(begin, sizes[i] * mask)
  return [None] + out_grads
Ejemplo n.º 9
0
  def testShapeFunctionEdgeCases(self):
    # Non-vector dimensions.
    with self.assertRaises(errors_impl.InvalidArgumentError):
      array_ops.fill([[0, 1], [2, 3]], 1.0)

    # Non-scalar value.
    with self.assertRaises(errors_impl.InvalidArgumentError):
      array_ops.fill([3, 2], [1.0, 2.0])
Ejemplo n.º 10
0
def _SegmentMeanGrad(op, grad):
    """Gradient for SegmentMean."""
    input_rank = array_ops.rank(op.inputs[0])
    ones_shape = array_ops.concat(
        0, [array_ops.shape(op.inputs[1]), array_ops.fill(array_ops.expand_dims(input_rank - 1, 0), 1)]
    )
    ones = array_ops.fill(ones_shape, constant_op.constant(1, dtype=grad.dtype))
    scaled_grad = grad * math_ops.inv(math_ops.segment_sum(ones, op.inputs[1]))
    return array_ops.gather(scaled_grad, op.inputs[1]), None
Ejemplo n.º 11
0
  def testFillNegative(self):
    with self.test_session():
      for shape in (-1,), (2, -1), (-1, 2), (-2), (-3):
        with self.assertRaises(ValueError):
          array_ops.fill(shape, 7)

      # Using a placeholder so this won't be caught in static analysis.
      dims = array_ops.placeholder(dtypes_lib.int32)
      fill_t = array_ops.fill(dims, 3.0)
      for shape in (-1,), (2, -1), (-1, 2), (-2), (-3):
        with self.assertRaises(errors_impl.InvalidArgumentError):
          fill_t.eval({dims: shape})
Ejemplo n.º 12
0
  def testAssignNonStrictShapeChecking(self):
    with self.cached_session():
      data = array_ops.fill([1024, 1024], 0)
      p = variables.VariableV1([1])
      a = state_ops.assign(p, data, validate_shape=False)
      a.op.run()
      self.assertAllEqual(p.eval(), data.eval())

      # Assign to yet another shape
      data2 = array_ops.fill([10, 10], 1)
      a2 = state_ops.assign(p, data2, validate_shape=False)
      a2.op.run()
      self.assertAllEqual(p.eval(), data2.eval())
Ejemplo n.º 13
0
 def _CreateDenseMaskAndBegin(sizes, concat_dim):
   """Create variables for iteratively slicing a dense gradients tensor."""
   # Since shape is 1-D, shape_of_shape = [rank-of-inputs]
   shape_of_shape = array_ops.shape(sizes[0])
   # Make a vector of length equal to the input's dimensions,
   # with 0's everywhere and 1 in the concat dim position.
   # Note: Can't use sparse_to_dense since it isn't GPU-capable (for now)
   mask = array_ops.concat([
       array_ops.fill(array_ops.expand_dims(concat_dim, 0), 0), [1],
       array_ops.fill(shape_of_shape - concat_dim - 1, 0)
   ], 0)
   begin = array_ops.fill(shape_of_shape, 0)
   return mask, begin
  def testDenseToSparseBatchDatasetWithUnknownShape(self):
    components = np.random.randint(5, size=(40,)).astype(np.int32)
    iterator = (
        dataset_ops.Dataset.from_tensor_slices(components)
        .map(lambda x: array_ops.fill([x, x], x)).apply(
            batching.dense_to_sparse_batch(
                4, [5, None])).make_initializable_iterator())
    init_op = iterator.initializer
    get_next = iterator.get_next()

    with self.cached_session() as sess:
      self.evaluate(init_op)

      for start in range(0, len(components), 4):
        results = self.evaluate(get_next)
        self.assertAllEqual([[i, j, z]
                             for i, c in enumerate(components[start:start + 4])
                             for j in range(c)
                             for z in range(c)], results.indices)
        self.assertAllEqual([
            c
            for c in components[start:start + 4] for _ in range(c)
            for _ in range(c)
        ], results.values)
        self.assertAllEqual([
            min(4,
                len(components) - start), 5,
            np.max(components[start:start + 4])
        ], results.dense_shape)

      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)
Ejemplo n.º 15
0
 def testDtype(self):
   with self.test_session():
     d = array_ops.fill([2, 3], 12., name="fill")
     self.assertEqual(d.get_shape(), [2, 3])
     # Test default type for both constant size and dynamic size
     z = array_ops.zeros([2, 3])
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.zeros([2, 3]))
     z = array_ops.zeros(array_ops.shape(d))
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.zeros([2, 3]))
     # Test explicit type control
     for dtype in [
         dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int32,
         dtypes_lib.uint8, dtypes_lib.int16, dtypes_lib.int8,
         dtypes_lib.complex64, dtypes_lib.complex128, dtypes_lib.int64,
         dtypes_lib.bool, dtypes_lib.string
     ]:
       z = array_ops.zeros([2, 3], dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       z_value = z.eval()
       self.assertFalse(np.any(z_value))
       self.assertEqual((2, 3), z_value.shape)
       z = array_ops.zeros(array_ops.shape(d), dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       z_value = z.eval()
       self.assertFalse(np.any(z_value))
       self.assertEqual((2, 3), z_value.shape)
Ejemplo n.º 16
0
    def body(time, outputs_ta, state, inputs, finished, sequence_lengths):
      """Internal while_loop body.

      Args:
        time: scalar int32 tensor.
        outputs_ta: structure of TensorArray.
        state: (structure of) state tensors and TensorArrays.
        inputs: (structure of) input tensors.
        finished: bool tensor (keeping track of what's finished).
        sequence_lengths: int32 tensor (keeping track of time of finish).

      Returns:
        `(time + 1, outputs_ta, next_state, next_inputs, next_finished,
          next_sequence_lengths)`.
        ```
      """
      (next_outputs, decoder_state, next_inputs,
       decoder_finished) = decoder.step(time, inputs, state)
      next_finished = math_ops.logical_or(decoder_finished, finished)
      if maximum_iterations is not None:
        next_finished = math_ops.logical_or(
            next_finished, time + 1 >= maximum_iterations)
      next_sequence_lengths = array_ops.where(
          math_ops.logical_and(math_ops.logical_not(finished), next_finished),
          array_ops.fill(array_ops.shape(sequence_lengths), time + 1),
          sequence_lengths)

      nest.assert_same_structure(state, decoder_state)
      nest.assert_same_structure(outputs_ta, next_outputs)
      nest.assert_same_structure(inputs, next_inputs)

      # Zero out output values past finish
      if impute_finished:
        emit = nest.map_structure(
            lambda out, zero: array_ops.where(finished, zero, out),
            next_outputs,
            zero_outputs)
      else:
        emit = next_outputs

      # Copy through states past finish
      def _maybe_copy_state(new, cur):
        # TensorArrays and scalar states get passed through.
        if isinstance(cur, tensor_array_ops.TensorArray):
          pass_through = True
        else:
          new.set_shape(cur.shape)
          pass_through = (new.shape.ndims == 0)
        return new if pass_through else array_ops.where(finished, cur, new)

      if impute_finished:
        next_state = nest.map_structure(
            _maybe_copy_state, decoder_state, state)
      else:
        next_state = decoder_state

      outputs_ta = nest.map_structure(lambda ta, out: ta.write(time, out),
                                      outputs_ta, emit)
      return (time + 1, outputs_ta, next_state, next_inputs, next_finished,
              next_sequence_lengths)
Ejemplo n.º 17
0
 def testDtype(self):
   with self.test_session():
     d = array_ops.fill([2, 3], 12., name="fill")
     self.assertEqual(d.get_shape(), [2, 3])
     # Test default type for both constant size and dynamic size
     z = array_ops.ones([2, 3])
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     z = array_ops.ones(array_ops.shape(d))
     self.assertEqual(z.dtype, dtypes_lib.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     # Test explicit type control
     for dtype in (dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int32,
                   dtypes_lib.uint8, dtypes_lib.int16, dtypes_lib.int8,
                   dtypes_lib.complex64, dtypes_lib.complex128,
                   dtypes_lib.int64, dtypes_lib.bool):
       z = array_ops.ones([2, 3], dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))
       z = array_ops.ones(array_ops.shape(d), dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))
Ejemplo n.º 18
0
def _DefaultGradYs(grad_ys, ys, colocate_gradients_with_ops):
  """Fill in default values for grad_ys.

  Args:
    grad_ys: List of gradients, can contain None.
    ys: List of tensors.
    colocate_gradients_with_ops: If True, try colocating gradients with
      the corresponding op.

  Returns:
    A list of gradients to use, without None.

  Raises:
    ValueError: If one of the grad_ys is invalid.
  """
  if len(grad_ys) != len(ys):
    raise ValueError("Passed %d grad_ys for %d ys" % (len(grad_ys), len(ys)))
  grad_ys = ops.convert_n_to_tensor_or_indexed_slices(grad_ys, name="grad_y")
  for i in xrange(len(grad_ys)):
    grad_y = grad_ys[i]
    y = ys[i]
    if grad_y is None:
      with _maybe_colocate_with(y.op, colocate_gradients_with_ops):
        grad_ys[i] = array_ops.fill(
            array_ops.shape(y), constant_op.constant(
                1, dtype=y.dtype))
    else:
      if grad_y.dtype != y.dtype:
        raise ValueError("Y and ys_grad must be of the same type, "
                         "not y: %s, ys_grad: %s " %
                         (dtypes.as_dtype(y.dtype).name,
                          dtypes.as_dtype(grad_y.dtype).name))
  return grad_ys
Ejemplo n.º 19
0
 def testLargeFetch(self):
   server = self._cached_server
   with session.Session(server.target, config=self._useRPCConfig()) as sess:
     c = array_ops.fill([10000, 3000], 0.5)
     expected_val = np.empty([10000, 3000], dtype=np.float32)
     expected_val.fill(0.5)
     self.assertAllEqual(expected_val, sess.run(c))
  def testParallelAssignWithoutLocking(self):
    # We need each thread to keep its own device stack or the device scopes
    # won't be properly nested.
    ops.get_default_graph().switch_to_thread_local()
    with self.cached_session() as sess:
      ones_t = array_ops.fill([1024, 1024], float(1))
      p = variables.Variable(array_ops.zeros([1024, 1024]))
      assigns = [
          state_ops.assign(p, math_ops.multiply(ones_t, float(i)), False)
          for i in range(1, 21)
      ]
      self.evaluate(variables.global_variables_initializer())

      def run_assign(assign_op):
        self.evaluate(assign_op)

      threads = [
          self.checkedThread(
              target=run_assign, args=(assign_op,)) for assign_op in assigns
      ]
      for t in threads:
        t.start()
      for t in threads:
        t.join()

      vals = self.evaluate(p)

      # Assert every element is taken from one of the assignments.
      self.assertTrue((vals > 0).all())
      self.assertTrue((vals <= 20).all())
Ejemplo n.º 21
0
 def _compare(self, dims, val, np_ans, use_gpu):
   ctx = context.get_default_context()
   device = "GPU:0" if (use_gpu and ctx.num_gpus()) else "CPU:0"
   with ops.device(device):
     tf_ans = array_ops.fill(dims, val, name="fill")
     out = tf_ans.numpy()
   self.assertAllClose(np_ans, out)
Ejemplo n.º 22
0
def _Solve(a, b, c):
    """Return solution of a quadratic minimization.

  The optimization equation is:
       f(a, b, c) = argmin_w{1/2 * a * w^2 + b * w + c * |w|}
  we get optimal solution w*:
       w* = -(b - sign(b)*c)/a if |b| > c else w* = 0

  REQUIRES: Dimensionality of a and b must be same

  Args:
    a: A Tensor
    b: A Tensor
    c: A Tensor with one element.

  Returns:
    A Tensor w, which is solution for the equation
  """
    with ops.name_scope("solve_" + b.op.name):
        c = ops.convert_to_tensor(c)
        k = array_ops.fill(array_ops.shape(b), c)
        zero_t = array_ops.zeros(array_ops.shape(b), dtype=b.dtype)
        w = (c * math_ops.sign(b) - b) / a
        w = math_ops.select(math_ops.less(math_ops.abs(b), k), zero_t, w)
        return w
Ejemplo n.º 23
0
  def _writeDummySavedModel(self, path, feature_name):
    """Writes a classifier with two input features to the given path."""
    with ops.Graph().as_default():
      examples = array_ops.placeholder(dtypes.string, name="input_node")
      feature_configs = {
          feature_name: parsing_ops.FixedLenFeature(shape=[],
                                                    dtype=dtypes.float32),
      }
      features = parsing_ops.parse_example(examples, feature_configs)
      feature = features[feature_name]

      variable_node = variables.VariableV1(1.0, name="variable_node")
      scores = math_ops.multiply(variable_node, feature, name="output_node")
      class_feature = array_ops.fill(array_ops.shape(feature),
                                     "class_%s" % feature_name)
      classes = array_ops.transpose(class_feature)

      with session.Session() as sess:
        sess.run(variables.global_variables_initializer())
        signature = (
            signature_def_utils.classification_signature_def(
                examples=examples,
                classes=classes,
                scores=scores,))
        builder = saved_model_builder.SavedModelBuilder(path)
        builder.add_meta_graph_and_variables(
            sess,
            [tag_constants.SERVING],
            signature_def_map={
                signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    signature,
            },)
        builder.save(as_text=True)
Ejemplo n.º 24
0
  def testConsumeWindowDatasetMoreThanOnce(self):
    components = np.random.randint(50, size=(200,)).astype(np.int64)

    def reduce_func(key, window):
      # Apply two different kinds of padding to the input: tight
      # padding, and quantized (to a multiple of 10) padding.
      return dataset_ops.Dataset.zip((window.padded_batch(
          4,
          padded_shapes=tensor_shape.TensorShape([None])), window.padded_batch(
              4, padded_shapes=ops.convert_to_tensor([(key + 1) * 10])),))

    iterator = dataset_ops.Iterator.from_dataset(
        dataset_ops.Dataset.from_tensor_slices(components)
        .map(lambda x: array_ops.fill([math_ops.cast(x, dtypes.int32)], x))
        .group_by_window(
            lambda x: math_ops.cast(array_ops.shape(x)[0] // 10, dtypes.int64),
            reduce_func, 4))
    init_op = iterator.initializer
    get_next = iterator.get_next()

    with self.test_session() as sess:
      sess.run(init_op)
      counts = []
      with self.assertRaises(errors.OutOfRangeError):
        while True:
          tight_result, multiple_of_10_result = sess.run(get_next)
          self.assertEqual(0, multiple_of_10_result.shape[1] % 10)
          self.assertAllEqual(tight_result,
                              multiple_of_10_result[:, :tight_result.shape[1]])
          counts.append(tight_result.shape[0])
      self.assertEqual(len(components), sum(counts))
Ejemplo n.º 25
0
 def _unbounded_exponential_log_prob(x):
   """An exponential distribution with log-likelihood NaN for x < 0."""
   per_element_potentials = array_ops.where(
       x < 0.,
       array_ops.fill(array_ops.shape(x), x.dtype.as_numpy_dtype(np.nan)),
       -x)
   return math_ops.reduce_sum(per_element_potentials)
Ejemplo n.º 26
0
def calculate_reshape(original_shape, new_shape, validate=False, name=None):
  """Calculates the reshaped dimensions (replacing up to one -1 in reshape)."""
  batch_shape_static = tensor_util.constant_value_as_shape(new_shape)
  if batch_shape_static.is_fully_defined():
    return np.int32(batch_shape_static.as_list()), batch_shape_static, []
  with ops.name_scope(name, "calculate_reshape", [original_shape, new_shape]):
    original_size = math_ops.reduce_prod(original_shape)
    implicit_dim = math_ops.equal(new_shape, -1)
    size_implicit_dim = (
        original_size // math_ops.maximum(1, -math_ops.reduce_prod(new_shape)))
    new_ndims = array_ops.shape(new_shape)
    expanded_new_shape = array_ops.where(  # Assumes exactly one `-1`.
        implicit_dim, array_ops.fill(new_ndims, size_implicit_dim), new_shape)
    validations = [] if not validate else [
        check_ops.assert_rank(
            original_shape, 1, message="Original shape must be a vector."),
        check_ops.assert_rank(
            new_shape, 1, message="New shape must be a vector."),
        check_ops.assert_less_equal(
            math_ops.count_nonzero(implicit_dim, dtype=dtypes.int32),
            1,
            message="At most one dimension can be unknown."),
        check_ops.assert_positive(
            expanded_new_shape, message="Shape elements must be >=-1."),
        check_ops.assert_equal(
            math_ops.reduce_prod(expanded_new_shape),
            original_size,
            message="Shape sizes do not match."),
    ]
    return expanded_new_shape, batch_shape_static, validations
Ejemplo n.º 27
0
  def quantiles_ready():
    """The subgraph for when the quantiles are ready."""
    quantized_feature = quantile_ops.quantiles([sparse_column_values], [],
                                               [quantile_buckets], [])
    quantized_feature = math_ops.cast(quantized_feature[0], dtypes.int64)
    quantized_feature = array_ops.reshape(quantized_feature, [-1])
    example_indices, _ = array_ops.split(
        sparse_column_indices, num_or_size_splits=2, axis=1)
    example_indices = array_ops.squeeze(example_indices, [1])
    filtered_gradients = array_ops.gather(gradients, example_indices)
    filtered_hessians = array_ops.gather(hessians, example_indices)
    filtered_partition_ids = array_ops.gather(example_partition_ids,
                                              example_indices)
    unique_partitions, mapped_partitions = array_ops.unique(
        example_partition_ids)

    # Compute aggregate stats for each partition.
    per_partition_gradients = math_ops.unsorted_segment_sum(
        gradients, mapped_partitions, array_ops.size(unique_partitions))
    per_partition_hessians = math_ops.unsorted_segment_sum(
        hessians, mapped_partitions, array_ops.size(unique_partitions))

    # Prepend a bias feature per partition that accumulates the stats for all
    # examples in that partition.
    bias_feature_ids = array_ops.fill(
        array_ops.shape(unique_partitions), _BIAS_FEATURE_ID)
    bias_feature_ids = math_ops.cast(bias_feature_ids, dtypes.int64)
    partition_ids = array_ops.concat(
        [unique_partitions, filtered_partition_ids], 0)
    filtered_gradients = array_ops.concat(
        [per_partition_gradients, filtered_gradients], 0)
    filtered_hessians = array_ops.concat(
        [per_partition_hessians, filtered_hessians], 0)
    bucket_ids = array_ops.concat([bias_feature_ids, quantized_feature], 0)
    return partition_ids, bucket_ids, filtered_gradients, filtered_hessians
Ejemplo n.º 28
0
def ndlstm_base_dynamic(inputs, noutput, scope=None, reverse=False):
  """Run an LSTM, either forward or backward.

  This is a 1D LSTM implementation using dynamic_rnn and
  the TensorFlow LSTM op.

  Args:
    inputs: input sequence (length, batch_size, ninput)
    noutput: depth of output
    scope: optional scope name
    reverse: run LSTM in reverse

  Returns:
    Output sequence (length, batch_size, noutput)
  """
  with variable_scope.variable_scope(scope, "SeqLstm", [inputs]):
    # TODO(tmb) make batch size, sequence_length dynamic
    # example: sequence_length = tf.shape(inputs)[0]
    _, batch_size, _ = _shape(inputs)
    lstm_cell = core_rnn_cell_impl.BasicLSTMCell(noutput, state_is_tuple=False)
    state = array_ops.zeros([batch_size, lstm_cell.state_size])
    sequence_length = int(inputs.get_shape()[0])
    sequence_lengths = math_ops.to_int64(
        array_ops.fill([batch_size], sequence_length))
    if reverse:
      inputs = array_ops.reverse_v2(inputs, [0])
    outputs, _ = rnn.dynamic_rnn(
        lstm_cell, inputs, sequence_lengths, state, time_major=True)
    if reverse:
      outputs = array_ops.reverse_v2(outputs, [0])
    return outputs
Ejemplo n.º 29
0
 def testPaddedBatchShortPadding(self):
   dataset = (
       dataset_ops.Dataset.from_tensor_slices(
           [6, 5, 5, 5, 5]).map(lambda x: array_ops.fill([x], x)).padded_batch(
               batch_size=4, padded_shapes=[5]))
   self.assertDatasetProduces(
       dataset, expected_error=(errors.DataLossError, ''))
Ejemplo n.º 30
0
def _state_to_olabel_unique(labels, num_labels, states, unique):
  """Sum state log probs to ilabel log probs using unique label indices."""

  num_label_states = _get_dim(labels, 1) + 1
  label_states = states[:, :, 1:num_label_states]
  blank_states = states[:, :, num_label_states:]

  unique_y, unique_idx = unique
  mul_reduce = _sum_states(unique_idx, label_states)

  num_frames = states.shape[0]
  batch_size = states.shape[1]
  num_states = num_label_states - 1
  batch_state_major = array_ops.transpose(mul_reduce, perm=[1, 2, 0])
  batch_state_major = array_ops.reshape(
      batch_state_major, [batch_size * num_states, num_frames])
  batch_offset = math_ops.range(batch_size, dtype=unique_y.dtype) * num_labels
  indices = unique_y + array_ops.expand_dims(batch_offset, axis=-1)
  indices = array_ops.reshape(indices, [-1, 1])
  scatter = array_ops.scatter_nd(
      indices=indices,
      updates=batch_state_major,
      shape=[batch_size * num_labels, num_frames])
  scatter = array_ops.reshape(scatter, [batch_size, num_labels, num_frames])
  scatter = array_ops.where(
      math_ops.equal(scatter, 0.0),
      array_ops.fill(array_ops.shape(scatter), math_ops.log(0.0)),
      scatter)
  label_olabels = array_ops.transpose(scatter, [2, 0, 1])
  label_olabels = label_olabels[:, :, 1:]

  blank_olabels = math_ops.reduce_logsumexp(
      blank_states, axis=2, keepdims=True)

  return array_ops.concat([blank_olabels, label_olabels], axis=-1)
Ejemplo n.º 31
0
 def testInt64(self):
     # Fill requires the first input to be an int32 tensor.
     self.assertAllEqual([1.0, 1.0],
                         array_ops.fill(
                             constant_op.constant([2], dtype=dtypes.int64),
                             constant_op.constant(1)))
Ejemplo n.º 32
0
def sparse_fill_empty_rows(sp_input, default_value, name=None):
  """Fills empty rows in the input 2-D `SparseTensor` with a default value.

  This op adds entries with the specified `default_value` at index
  `[row, 0]` for any row in the input that does not already have a value.

  For example, suppose `sp_input` has shape `[5, 6]` and non-empty values:

      [0, 1]: a
      [0, 3]: b
      [2, 0]: c
      [3, 1]: d

  Rows 1 and 4 are empty, so the output will be of shape `[5, 6]` with values:

      [0, 1]: a
      [0, 3]: b
      [1, 0]: default_value
      [2, 0]: c
      [3, 1]: d
      [4, 0]: default_value

  Note that the input may have empty columns at the end, with no effect on
  this op.

  The output `SparseTensor` will be in row-major order and will have the
  same shape as the input.

  This op also returns an indicator vector such that

      empty_row_indicator[i] = True iff row i was an empty row.

  Args:
    sp_input: A `SparseTensor` with shape `[N, M]`.
    default_value: The value to fill for empty rows, with the same type as
      `sp_input.`
    name: A name prefix for the returned tensors (optional)

  Returns:
    sp_ordered_output: A `SparseTensor` with shape `[N, M]`, and with all empty
      rows filled in with `default_value`.
    empty_row_indicator: A bool vector of length `N` indicating whether each
      input row was empty.

  Raises:
    TypeError: If `sp_input` is not a `SparseTensor`.
  """
  if not isinstance(sp_input, ops.SparseTensor):
    raise TypeError("Input must be a SparseTensor")

  with ops.op_scope([sp_input], name, "SparseFillEmptyRows"):
    default_value = ops.convert_to_tensor(default_value,
                                          dtype=sp_input.values.dtype)

    num_rows = math_ops.cast(sp_input.shape[0], dtypes.int32)
    all_row_indices = math_ops.cast(math_ops.range(num_rows), dtypes.int64)
    empty_row_indices, _ = array_ops.list_diff(all_row_indices,
                                               sp_input.indices[:, 0])
    empty_row_indicator = sparse_to_dense(
        empty_row_indices, array_ops.expand_dims(sp_input.shape[0], -1), True,
        False)

    empty_row_indices_as_column = array_ops.reshape(empty_row_indices, [-1, 1])
    additional_indices = array_ops.concat(
        1, [empty_row_indices_as_column,
            array_ops.zeros_like(empty_row_indices_as_column)])
    additional_values = array_ops.fill(
        array_ops.shape(empty_row_indices), default_value)

    all_indices_unordered = array_ops.concat(0, [sp_input.indices,
                                                 additional_indices])
    all_values_unordered = array_ops.concat(0, [sp_input.values,
                                                additional_values])
    sp_unordered_output = ops.SparseTensor(all_indices_unordered,
                                           all_values_unordered, sp_input.shape)
    sp_ordered_output = sparse_reorder(sp_unordered_output)

    return sp_ordered_output, empty_row_indicator
Ejemplo n.º 33
0
 def fill(x):
   return array_ops.fill([x], x)
def embedding_lookup_sparse(
    params,
    sp_ids,
    sp_weights,
    partition_strategy=None,  # no used
    name="embedding_lookup_sparse",
    combiner="mean",
    max_norm=None,
    return_trainable=False,
):
  """Provides a dynamic version of embedding_lookup_sparse
      similar with tf.nn.embedding_lookup_sparse.

    This op assumes that there is at least one id for each row in the dense tensor
    represented by sp_ids (i.e. there are no rows with empty features), and that
    all the indices of sp_ids are in canonical row-major order.

    It also assumes that all id values lie in the range [0, p0), where p0
    is the sum of the size of params along dimension 0.

    Args:
      params: A single `dynamic_embedding.Variable` instance representing
        the complete embedding tensor.
      sp_ids: N x M `SparseTensor` of int64 ids where N is typically batch size
        and M is arbitrary.
      sp_weights: either a `SparseTensor` of float / double weights, or `None` to
        indicate all weights should be taken to be 1. If specified, `sp_weights`
        must have exactly the same shape and indices as `sp_ids`.
      partition_strategy: No used.
      name: a name for the operation. Name is optional in graph mode and required
        in eager mode.
      combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
        and "sum" are supported. "sum" computes the weighted sum of the embedding
        results for each row. "mean" is the weighted sum divided by the total
        weight. "sqrtn" is the weighted sum divided by the square root of the sum
        of the squares of the weights.
      max_norm: If not `None`, each embedding is clipped if its l2-norm is larger
        than this value, before combining.
      return_trainable: optional, If True, also return TrainableWrapper create by
        `dynamic_embedding.embedding_lookup`

    Returns:
      combined_embeddings: A dense tensor representing the combined embeddings
        for the sparse ids. For each row in the dense tensor represented by
        `sp_ids`, the op looks up the embeddings for all ids in that row,
        multiplies them by the corresponding weight, and combines these embeddings
        as specified.

        In other words, if

          `shape(combined params) = [+infinity, dim]`

        and

          `shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn]`

        then

          `shape(output) = [d0, dim]`.

        For instance, if params dim=20, and sp_ids / sp_weights are

          ```python
          [0, 0]: id 1, weight 2.0
          [0, 1]: id 3, weight 0.5
          [1, 0]: id 0, weight 1.0
          [2, 3]: id 1, weight 3.0
          ```

        with `combiner`="mean", then the output will be a 3x20 matrix where

          ```python
          output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
          output[1, :] = (params[0, :] * 1.0) / 1.0
          output[2, :] = (params[1, :] * 3.0) / 3.0
          ```
      trainable_wrap:
        A TrainableWrapper object used to fill the Optimizers `var_list`
          Only provided if `return_trainable` is True.
    Raises:
      TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is
        neither `None` nor `SparseTensor`.
      ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
    """
  if combiner not in ("mean", "sqrtn", "sum"):
    raise ValueError("combiner must be one of 'mean', 'sqrtn' or 'sum'")

  if not isinstance(sp_ids, sparse_tensor.SparseTensor):
    raise TypeError("sp_ids must be SparseTensor")

  ignore_weights = sp_weights is None
  if not ignore_weights:
    if not isinstance(sp_weights, sparse_tensor.SparseTensor):
      raise TypeError("sp_weights must be either None or SparseTensor")

  scope = variable_scope.get_variable_scope()
  full_name = scope.name + "/" + name if scope.name else name
  with ops.name_scope(full_name + "/"):
    segment_ids = sp_ids.indices[:, 0]
    if segment_ids.dtype != dtypes.int32:
      segment_ids = math_ops.cast(segment_ids, dtypes.int32)

    ids = sp_ids.values
    ids, idx = array_ops.unique(ids)

    embeddings, trainable_ = embedding_lookup(
        params,
        ids,
        name=name + "/embedding_lookup",
        partition_strategy=partition_strategy,
        max_norm=max_norm,
        return_trainable=True,
    )
    if embeddings.dtype in (dtypes.float16, dtypes.bfloat16):
      embeddings = math_ops.cast(embeddings, dtypes.float32)
    if not ignore_weights:
      weights = sp_weights.values
      if weights.dtype != embeddings.dtype:
        weights = math_ops.cast(weights, embeddings.dtype)

      embeddings = array_ops.gather(embeddings, idx)

      # Reshape weights to allow broadcast
      ones = array_ops.fill(
          array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
      bcast_weights_shape = array_ops.concat([array_ops.shape(weights), ones],
                                             0)

      orig_weights_shape = weights.get_shape()
      weights = array_ops.reshape(weights, bcast_weights_shape)

      # Set the weight shape, since after reshaping to bcast_weights_shape,
      # the shape becomes None.
      if embeddings.get_shape().ndims is not None:
        weights.set_shape(
            orig_weights_shape.concatenate(
                [1 for _ in range(embeddings.get_shape().ndims - 1)]))

      embeddings *= weights

      if combiner == "sum":
        embeddings = math_ops.segment_sum(embeddings, segment_ids, name=name)
      elif combiner == "mean":
        embeddings = math_ops.segment_sum(embeddings, segment_ids)
        weight_sum = math_ops.segment_sum(weights, segment_ids)
        embeddings = math_ops.div(embeddings, weight_sum, name=name)
      elif combiner == "sqrtn":
        embeddings = math_ops.segment_sum(embeddings, segment_ids)
        weights_squared = math_ops.pow(weights, 2)
        weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
        weight_sum_sqrt = math_ops.sqrt(weight_sum)
        embeddings = math_ops.div(embeddings, weight_sum_sqrt, name=name)
      else:
        assert False, "Unrecognized combiner"
    else:
      assert idx is not None
      if combiner == "sum":
        embeddings = de.math.sparse_segment_sum(embeddings,
                                                idx,
                                                segment_ids,
                                                name=name)
      elif combiner == "mean":
        embeddings = math_ops.sparse_segment_mean(embeddings,
                                                  idx,
                                                  segment_ids,
                                                  name=name)
      elif combiner == "sqrtn":
        embeddings = math_ops.sparse_segment_sqrt_n(embeddings,
                                                    idx,
                                                    segment_ids,
                                                    name=name)
      else:
        assert False, "Unrecognized combiner"

    return (embeddings, trainable_) if return_trainable else embeddings
 def testArrayFill(self):
     array_ops.fill(constant_op.constant([2], dtype=dtypes.int64),
                    constant_op.constant(1))
Ejemplo n.º 36
0
def _DefaultGradYs(grad_ys,
                   ys,
                   colocate_gradients_with_ops,
                   gradient_uid="__unsupported__"):
    """Fill in default values for grad_ys.

  Args:
    grad_ys: List of gradients, can contain None.
    ys: List of tensors.
    colocate_gradients_with_ops: If True, try colocating gradients with
      the corresponding op.
    gradient_uid: A unique identifier within the graph indicating
      which invocation of gradients is being executed. Used to cluster
      ops for compilation.

  Returns:
    A list of gradients to use, without None.

  Raises:
    ValueError: If sizes of gradients and inputs don't match
    TypeError: If type of any gradient is not valid for its input.
  """
    if len(grad_ys) != len(ys):
        raise ValueError("Passed %d grad_ys for %d ys" %
                         (len(grad_ys), len(ys)))
    grad_ys = ops.convert_n_to_tensor_or_indexed_slices(grad_ys, name="grad_y")
    new_grad_ys = []
    for i, (y, grad_y) in enumerate(zip(ys, grad_ys)):
        with _maybe_colocate_with(y.op, gradient_uid,
                                  colocate_gradients_with_ops):
            if grad_y is None:
                if y.dtype.is_complex:
                    raise TypeError(
                        "Gradients of complex tensors must set grad_ys (y.dtype = %r)"
                        % y.dtype)
                new_grad_ys.append(
                    array_ops.fill(
                        array_ops.shape(y),
                        constant_op.constant(1,
                                             dtype=y.dtype,
                                             name="grad_ys_%d" % i)))
                continue
            if y.dtype.is_floating or y.dtype.is_integer:
                if not grad_y.dtype.is_floating and not grad_y.dtype.is_integer:
                    raise TypeError(
                        "Gradient type %s generated for real or "
                        "integer-valued tensor %s with type %s must be "
                        "real or integer" %
                        (dtypes.as_dtype(grad_y.dtype).name, y,
                         dtypes.as_dtype(y.dtype).name))
            elif y.dtype.is_complex:
                if not grad_y.dtype.is_complex:
                    raise TypeError(
                        "Gradient type %s generated for complex-valued "
                        "tensor %s with type %s must be real" %
                        (dtypes.as_dtype(grad_y.dtype).name, y,
                         dtypes.as_dtype(y.dtype).name))
            elif y.dtype == dtypes.variant:
                if grad_y.dtype != dtypes.variant:
                    raise TypeError("Gradient type %s generated for variant "
                                    "tensor %s with type %s must be variant" %
                                    (dtypes.as_dtype(grad_y.dtype).name, y,
                                     dtypes.as_dtype(y.dtype).name))
            elif y.dtype == dtypes.resource:
                # We assume y is the handle of a ResourceVariable. The gradient of a
                # ResourceVariable should be a numeric value, not another resource.
                if grad_y.dtype == dtypes.resource:
                    raise TypeError(
                        "Input gradient %s for resource tensor %s should not "
                        "be a resource" % (grad_y, y))
            else:
                raise TypeError("Tensor %s with type %s must be numeric "
                                "to obtain a default gradient" %
                                (y, dtypes.as_dtype(y.dtype).name))
            # Create a grad_y tensor in the name scope of the gradient.
            # Required for TensorArrays to identify which gradient call a
            # grad_y value is coming from.
            if isinstance(grad_y, ops.IndexedSlices):
                new_grad_ys.append(
                    ops.IndexedSlices(
                        indices=(array_ops.identity(
                            grad_y.indices, name="grad_ys_%d_indices" %
                            i) if isinstance(grad_y.indices, ops.Tensor) else
                                 grad_y.indices),
                        values=(array_ops.identity(
                            grad_y.values, name="grad_ys_%d_values" %
                            i) if isinstance(grad_y.values, ops.Tensor) else
                                grad_y.values),
                        dense_shape=(array_ops.identity(
                            grad_y.dense_shape, name="grad_ys_%d_shape" %
                            i) if isinstance(grad_y.dense_shape, ops.Tensor)
                                     else grad_y.dense_shape)))
            else:
                new_grad_ys.append(
                    array_ops.identity(grad_y, name="grad_ys_%d" % i))

    return new_grad_ys
Ejemplo n.º 37
0
def _fast_fill(value, shape, dtype):
    return array_ops.fill(constant_op.constant(shape, dtype=dtypes.int32),
                          constant_op.constant(value, dtype=dtype))
Ejemplo n.º 38
0
 def testPaddedBatchEmptyTensors(self):
   dataset = (
       dataset_ops.Dataset.from_tensor_slices(
           [0, 0, 0, 0]).map(lambda x: array_ops.fill([x], x)).padded_batch(
               batch_size=4, padded_shapes=[-1]))
   self.assertDatasetProduces(dataset, expected_output=[[[], [], [], []]])
Ejemplo n.º 39
0
 def fill_tuple(x):
   filled = array_ops.fill([x], x)
   return (filled, string_ops.as_string(filled))
Ejemplo n.º 40
0
 def build_dataset(seq_lens):
   return dataset_ops.Dataset.from_tensor_slices(seq_lens).map(
       lambda x: array_ops.fill([x], x)).padded_batch(
           4, padded_shapes=[-1])
    def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        """Apply gradients to variables.

    This contains most of the synchronization implementation and also wraps the
    apply_gradients() from the real optimizer.

    Args:
      grads_and_vars: List of (gradient, variable) pairs as returned by
        compute_gradients().
      global_step: Optional Variable to increment by one after the
        variables have been updated.
      name: Optional name for the returned operation.  Default to the
        name passed to the Optimizer constructor.

    Returns:
      train_op: The op to dequeue a token so the replicas can exit this batch
      and start the next one. This is executed by each replica.

    Raises:
      ValueError: If the grads_and_vars is empty.
      ValueError: If global step is not provided, the staleness cannot be
        checked.
    """
        if not grads_and_vars:
            raise ValueError("Must supply at least one variable")

        if global_step is None:
            raise ValueError("Global step is required to check staleness")

        self._global_step = global_step
        train_ops = []
        aggregated_grad = []
        var_list = []

        # local_anchor op will be placed on this worker task by default.
        local_anchor = control_flow_ops.no_op()
        # Colocating local_step variable prevents it being placed on the PS.
        with ops.colocate_with(local_anchor):
            self._local_step = variable_scope.variable(
                initial_value=0,
                trainable=False,
                collections=[ops.GraphKeys.LOCAL_VARIABLES],
                dtype=global_step.dtype.base_dtype,
                name="sync_rep_local_step")

        self.local_step_init_op = state_ops.assign(self._local_step,
                                                   global_step)
        chief_init_ops = [self.local_step_init_op]
        self.ready_for_local_init_op = variables.report_uninitialized_variables(
            variables.global_variables())

        with ops.name_scope(None, self._name):
            for grad, var in grads_and_vars:
                var_list.append(var)
                with ops.device(var.device):
                    # Dense gradients.
                    if grad is None:
                        aggregated_grad.append(None)  # pass-through.
                        continue
                    elif isinstance(grad, ops.Tensor):
                        grad_accum = data_flow_ops.ConditionalAccumulator(
                            grad.dtype,
                            shape=var.get_shape(),
                            shared_name=var.name + "/grad_accum")
                        train_ops.append(
                            grad_accum.apply_grad(grad,
                                                  local_step=self._local_step))
                        aggregated_grad.append(
                            grad_accum.take_grad(self._replicas_to_aggregate))
                    else:
                        if not isinstance(grad, ops.IndexedSlices):
                            raise ValueError("Unknown grad type!")
                        grad_accum = data_flow_ops.SparseConditionalAccumulator(
                            grad.dtype,
                            shape=(),
                            shared_name=var.name + "/grad_accum")
                        train_ops.append(
                            grad_accum.apply_indexed_slices_grad(
                                grad, local_step=self._local_step))
                        aggregated_grad.append(
                            grad_accum.take_indexed_slices_grad(
                                self._replicas_to_aggregate))

                    self._accumulator_list.append((grad_accum, var.device))

            aggregated_grads_and_vars = zip(aggregated_grad, var_list)

            # sync_op will be assigned to the same device as the global step.
            with ops.device(global_step.device), ops.name_scope(""):
                update_op = self._opt.apply_gradients(
                    aggregated_grads_and_vars, global_step)

            # Create token queue.
            with ops.device(global_step.device), ops.name_scope(""):
                sync_token_queue = (data_flow_ops.FIFOQueue(
                    -1,
                    global_step.dtype.base_dtype,
                    shapes=(),
                    name="sync_token_q",
                    shared_name="sync_token_q"))
                self._sync_token_queue = sync_token_queue

                # dummy_queue is passed to the queue runner. Don't use the real queues
                # because the queue runner doesn't automatically reopen it once it
                # closed queues in PS devices.
                dummy_queue = (data_flow_ops.FIFOQueue(
                    1,
                    types_pb2.DT_INT32,
                    shapes=(),
                    name="dummy_queue",
                    shared_name="dummy_queue"))

            with ops.device(global_step.device), ops.name_scope(""):
                # Replicas have to wait until they can get a token from the token queue.
                with ops.control_dependencies(train_ops):
                    token = sync_token_queue.dequeue()
                train_op = state_ops.assign(self._local_step, token)

                with ops.control_dependencies([update_op]):
                    # Sync_op needs to insert tokens to the token queue at the end of the
                    # step so the replicas can fetch them to start the next step.
                    tokens = array_ops.fill([self._tokens_per_step],
                                            global_step)
                    sync_op = sync_token_queue.enqueue_many((tokens, ))

                if self._variable_averages is not None:
                    with ops.control_dependencies([sync_op
                                                   ]), ops.name_scope(""):
                        sync_op = self._variable_averages.apply(
                            self._variables_to_average)

                self._chief_queue_runner = queue_runner.QueueRunner(
                    dummy_queue, [sync_op])
            for accum, dev in self._accumulator_list:
                with ops.device(dev):
                    chief_init_ops.append(
                        accum.set_global_step(global_step,
                                              name="SetGlobalStep"))
            self.chief_init_op = control_flow_ops.group(*(chief_init_ops))
            self._gradients_applied = True
            return train_op
Ejemplo n.º 42
0
def kernel_classifier_distance_and_std_from_activations(
        real_activations,
        generated_activations,
        max_block_size=10,
        dtype=None):
    """Kernel "classifier" distance for evaluating a generative model.

	This methods computes the kernel classifier distance from activations of
	real images and generated images. This can be used independently of the
	kernel_classifier_distance() method, especially in the case of using large
	batches during evaluation where we would like to precompute all of the
	activations before computing the classifier distance, or if we want to
	compute multiple metrics based on the same images. It also returns a rough
	estimate of the standard error of the estimator.

	This technique is described in detail in https://arxiv.org/abs/1801.01401.
	Given two distributions P and Q of activations, this function calculates

		E_{X, X' ~ P}[k(X, X')] + E_{Y, Y' ~ Q}[k(Y, Y')]
		  - 2 E_{X ~ P, Y ~ Q}[k(X, Y)]

	where k is the polynomial kernel

		k(x, y) = ( x^T y / dimension + 1 )^3.

	This captures how different the distributions of real and generated images'
	visual features are. Like the Frechet distance (and unlike the Inception
	score), this is a true distance and incorporates information about the
	target images. Unlike the Frechet score, this function computes an
	*unbiased* and asymptotically normal estimator, which makes comparing
	estimates across models much more intuitive.

	The estimator used takes time quadratic in max_block_size. Larger values of
	max_block_size will decrease the variance of the estimator but increase the
	computational cost. This differs slightly from the estimator used by the
	original paper; it is the block estimator of https://arxiv.org/abs/1307.1954.
	The estimate of the standard error will also be more reliable when there are
	more blocks, i.e. when max_block_size is smaller.

	NOTE: the blocking code assumes that real_activations and
	generated_activations are both in random order. If either is sorted in a
	meaningful order, the estimator will behave poorly.

	Args:
	  real_activations: 2D Tensor containing activations of real data. Shape is
		[batch_size, activation_size].
	  generated_activations: 2D Tensor containing activations of generated data.
		Shape is [batch_size, activation_size].
	  max_block_size: integer, default 1024. The distance estimator splits samples
		into blocks for computational efficiency. Larger values are more
		computationally expensive but decrease the variance of the distance
		estimate. Having a smaller block size also gives a better estimate of the
		standard error.
	  dtype: if not None, coerce activations to this dtype before computations.

	Returns:
	 The Kernel Inception Distance. A floating-point scalar of the same type
	   as the output of the activations.
	 An estimate of the standard error of the distance estimator (a scalar of
	   the same type).
	"""

    real_activations.shape.assert_has_rank(2)
    generated_activations.shape.assert_has_rank(2)
    real_activations.shape[1].assert_is_compatible_with(
        generated_activations.shape[1])

    if dtype is None:
        dtype = real_activations.dtype
        assert generated_activations.dtype == dtype
    else:
        real_activations = math_ops.cast(real_activations, dtype)
        generated_activations = math_ops.cast(generated_activations, dtype)

    # Figure out how to split the activations into blocks of approximately
    # equal size, with none larger than max_block_size.
    n_r = array_ops.shape(real_activations)[0]
    n_g = array_ops.shape(generated_activations)[0]

    n_bigger = math_ops.maximum(n_r, n_g)
    n_blocks = math_ops.to_int32(math_ops.ceil(n_bigger / max_block_size))

    v_r = n_r // n_blocks
    v_g = n_g // n_blocks

    n_plusone_r = n_r - v_r * n_blocks
    n_plusone_g = n_g - v_g * n_blocks

    sizes_r = array_ops.concat([
        array_ops.fill([n_blocks - n_plusone_r], v_r),
        array_ops.fill([n_plusone_r], v_r + 1),
    ], 0)
    sizes_g = array_ops.concat([
        array_ops.fill([n_blocks - n_plusone_g], v_g),
        array_ops.fill([n_plusone_g], v_g + 1),
    ], 0)

    zero = array_ops.zeros([1], dtype=dtypes.int32)
    inds_r = array_ops.concat([zero, math_ops.cumsum(sizes_r)], 0)
    inds_g = array_ops.concat([zero, math_ops.cumsum(sizes_g)], 0)

    dim = math_ops.cast(tf.shape(real_activations)[1], dtype)

    def compute_kid_block(i):
        'Compute the ith block of the KID estimate.'
        r_s = inds_r[i]
        r_e = inds_r[i + 1]
        r = real_activations[r_s:r_e]
        m = math_ops.cast(r_e - r_s, dtype)

        g_s = inds_g[i]
        g_e = inds_g[i + 1]
        g = generated_activations[g_s:g_e]
        n = math_ops.cast(g_e - g_s, dtype)

        k_rr = (math_ops.matmul(r, r, transpose_b=True) / dim + 1)**3
        k_rg = (math_ops.matmul(r, g, transpose_b=True) / dim + 1)**3
        k_gg = (math_ops.matmul(g, g, transpose_b=True) / dim + 1)**3
        return (-2 * math_ops.reduce_mean(k_rg) +
                (math_ops.reduce_sum(k_rr) - math_ops.trace(k_rr)) /
                (m * (m - 1)) +
                (math_ops.reduce_sum(k_gg) - math_ops.trace(k_gg)) / (n *
                                                                      (n - 1)))

    ests = functional_ops.map_fn(compute_kid_block,
                                 math_ops.range(n_blocks),
                                 dtype=dtype,
                                 back_prop=False)

    mn = math_ops.reduce_mean(ests)

    # nn_impl.moments doesn't use the Bessel correction, which we want here
    n_blocks_ = math_ops.cast(n_blocks, dtype)
    var = control_flow_ops.cond(
        math_ops.less_equal(n_blocks, 1),
        lambda: array_ops.constant(float('nan'), dtype=dtype),
        lambda: math_ops.reduce_sum(math_ops.square(ests - mn)) /
        (n_blocks_ - 1))

    return mn, math_ops.sqrt(var / n_blocks_)
Ejemplo n.º 43
0
def embedding_lookup_sparse(params,
                            sp_ids,
                            sp_weights,
                            partition_strategy="mod",
                            name=None,
                            combiner=None,
                            max_norm=None):
    """Computes embeddings for the given ids and weights.

  This op assumes that there is at least one id for each row in the dense tensor
  represented by sp_ids (i.e. there are no rows with empty features), and that
  all the indices of sp_ids are in canonical row-major order.

  It also assumes that all id values lie in the range [0, p0), where p0
  is the sum of the size of params along dimension 0.

  Args:
    params: A single tensor representing the complete embedding tensor,
      or a list of P tensors all of same shape except for the first dimension,
      representing sharded embedding tensors.  Alternatively, a
      `PartitionedVariable`, created by partitioning along dimension 0. Each
      element must be appropriately sized for the given `partition_strategy`.
    sp_ids: N x M `SparseTensor` of int64 ids where N is typically batch size
      and M is arbitrary.
    sp_weights: either a `SparseTensor` of float / double weights, or `None` to
      indicate all weights should be taken to be 1. If specified, `sp_weights`
      must have exactly the same shape and indices as `sp_ids`.
    partition_strategy: A string specifying the partitioning strategy, relevant
      if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
      is `"mod"`. See `tf.nn.embedding_lookup` for more details.
    name: Optional name for the op.
    combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
      and "sum" are supported.
      "sum" computes the weighted sum of the embedding results for each row.
      "mean" is the weighted sum divided by the total weight.
      "sqrtn" is the weighted sum divided by the square root of the sum of the
      squares of the weights.
    max_norm: If not `None`, each embedding is clipped if its l2-norm is
      larger than this value, before combining.

  Returns:
    A dense tensor representing the combined embeddings for the
    sparse ids. For each row in the dense tensor represented by `sp_ids`, the op
    looks up the embeddings for all ids in that row, multiplies them by the
    corresponding weight, and combines these embeddings as specified.

    In other words, if

      `shape(combined params) = [p0, p1, ..., pm]`

    and

      `shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn]`

    then

      `shape(output) = [d0, d1, ..., dn-1, p1, ..., pm]`.

    For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are

      ```python
      [0, 0]: id 1, weight 2.0
      [0, 1]: id 3, weight 0.5
      [1, 0]: id 0, weight 1.0
      [2, 3]: id 1, weight 3.0
      ```

    with `combiner`="mean", then the output will be a 3x20 matrix where

      ```python
      output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
      output[1, :] = (params[0, :] * 1.0) / 1.0
      output[2, :] = (params[1, :] * 3.0) / 3.0
      ```

  Raises:
    TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is
      neither `None` nor `SparseTensor`.
    ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
  """
    if combiner is None:
        logging.warn("The default value of combiner will change from \"mean\" "
                     "to \"sqrtn\" after 2016/11/01.")
        combiner = "mean"
    if combiner not in ("mean", "sqrtn", "sum"):
        raise ValueError("combiner must be one of 'mean', 'sqrtn' or 'sum'")
    if isinstance(params, variables.PartitionedVariable):
        params = list(params)  # Iterate to get the underlying Variables.
    if not isinstance(params, list):
        params = [params]
    if not isinstance(sp_ids, sparse_tensor.SparseTensor):
        raise TypeError("sp_ids must be SparseTensor")
    ignore_weights = sp_weights is None
    if not ignore_weights:
        if not isinstance(sp_weights, sparse_tensor.SparseTensor):
            raise TypeError("sp_weights must be either None or SparseTensor")
        sp_ids.values.get_shape().assert_is_compatible_with(
            sp_weights.values.get_shape())
        sp_ids.indices.get_shape().assert_is_compatible_with(
            sp_weights.indices.get_shape())
        sp_ids.dense_shape.get_shape().assert_is_compatible_with(
            sp_weights.dense_shape.get_shape())
        # TODO(yleon): Add enhanced node assertions to verify that sp_ids and
        # sp_weights have equal indices and shapes.

    with ops.name_scope(name, "embedding_lookup_sparse",
                        params + [sp_ids]) as name:
        segment_ids = sp_ids.indices[:, 0]
        if segment_ids.dtype != dtypes.int32:
            segment_ids = math_ops.cast(segment_ids, dtypes.int32)

        ids = sp_ids.values
        ids, idx = array_ops.unique(ids)

        embeddings = embedding_lookup(params,
                                      ids,
                                      partition_strategy=partition_strategy,
                                      max_norm=max_norm)
        if not ignore_weights:
            weights = sp_weights.values
            if weights.dtype != embeddings.dtype:
                weights = math_ops.cast(weights, embeddings.dtype)

            embeddings = array_ops.gather(embeddings, idx)

            # Reshape weights to allow broadcast
            ones = array_ops.fill(
                array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
            bcast_weights_shape = array_ops.concat(
                [array_ops.shape(weights), ones], 0)

            orig_weights_shape = weights.get_shape()
            weights = array_ops.reshape(weights, bcast_weights_shape)

            # Set the weight shape, since after reshaping to bcast_weights_shape,
            # the shape becomes None.
            if embeddings.get_shape().ndims is not None:
                weights.set_shape(
                    orig_weights_shape.concatenate(
                        [1 for _ in range(embeddings.get_shape().ndims - 1)]))

            embeddings *= weights

            if combiner == "sum":
                embeddings = math_ops.segment_sum(embeddings,
                                                  segment_ids,
                                                  name=name)
            elif combiner == "mean":
                embeddings = math_ops.segment_sum(embeddings, segment_ids)
                weight_sum = math_ops.segment_sum(weights, segment_ids)
                embeddings = math_ops.div(embeddings, weight_sum, name=name)
            elif combiner == "sqrtn":
                embeddings = math_ops.segment_sum(embeddings, segment_ids)
                weights_squared = math_ops.pow(weights, 2)
                weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
                weight_sum_sqrt = math_ops.sqrt(weight_sum)
                embeddings = math_ops.div(embeddings,
                                          weight_sum_sqrt,
                                          name=name)
            else:
                assert False, "Unrecognized combiner"
        else:
            assert idx is not None
            if combiner == "sum":
                embeddings = math_ops.sparse_segment_sum(embeddings,
                                                         idx,
                                                         segment_ids,
                                                         name=name)
            elif combiner == "mean":
                embeddings = math_ops.sparse_segment_mean(embeddings,
                                                          idx,
                                                          segment_ids,
                                                          name=name)
            elif combiner == "sqrtn":
                embeddings = math_ops.sparse_segment_sqrt_n(embeddings,
                                                            idx,
                                                            segment_ids,
                                                            name=name)
            else:
                assert False, "Unrecognized combiner"

        return embeddings
Ejemplo n.º 44
0
 def fill_tuple(x):
   return (x, array_ops.fill([x], x))
Ejemplo n.º 45
0
    def testPaddedBatchDataset(self):
        seq_lens = array_ops.placeholder(dtypes.int32, shape=[None])
        padded_shape = array_ops.placeholder(dtypes.int64, shape=[1])

        iterator = (dataset_ops.Dataset.from_tensor_slices(seq_lens).map(
            lambda x: array_ops.fill([x], x)).padded_batch(
                4, padded_shapes=padded_shape).make_initializable_iterator())

        init_op = iterator.initializer
        get_next = iterator.get_next()

        with self.test_session() as sess:
            # Test with random sequence lengths, and max padding.
            random_seq_lens = np.random.randint(20,
                                                size=(32, )).astype(np.int32)
            sess.run(init_op,
                     feed_dict={
                         padded_shape: [-1],
                         seq_lens: random_seq_lens
                     })
            for i in range(8):
                result = sess.run(get_next)
                padded_len = np.max(result)
                self.assertEqual((4, padded_len), result.shape)
                for j in range(4):
                    seq_len = random_seq_lens[(i * 4) + j]
                    self.assertAllEqual(result[j, :seq_len],
                                        [seq_len] * seq_len)
                    self.assertAllEqual(result[j, seq_len:],
                                        [0] * (padded_len - seq_len))
            with self.assertRaises(errors.OutOfRangeError):
                sess.run(get_next)

            # Test with random sequence lengths, and constant padding.
            sess.run(init_op,
                     feed_dict={
                         padded_shape: [25],
                         seq_lens: random_seq_lens
                     })
            for i in range(8):
                result = sess.run(get_next)
                self.assertEqual((4, 25), result.shape)
                for j in range(4):
                    seq_len = random_seq_lens[(i * 4) + j]
                    self.assertAllEqual(result[j, :seq_len],
                                        [seq_len] * seq_len)
                    self.assertAllEqual(result[j, seq_len:],
                                        [0] * (25 - seq_len))
            with self.assertRaises(errors.OutOfRangeError):
                sess.run(get_next)

            # Test correct handling of empty tensors.
            sess.run(init_op,
                     feed_dict={
                         padded_shape: [-1],
                         seq_lens: [0, 0, 0, 0]
                     })
            result = sess.run(get_next)
            self.assertAllEqual([[], [], [], []], result)
            with self.assertRaises(errors.OutOfRangeError):
                sess.run(get_next)

            # Test error handling with constant sequence lengths, and
            # too-short padding.
            sess.run(init_op,
                     feed_dict={
                         padded_shape: [5],
                         seq_lens: [6, 5, 5, 5]
                     })
            with self.assertRaises(errors.DataLossError):
                result = sess.run(get_next)
Ejemplo n.º 46
0
 def fill_tuple(x):
   filled = array_ops.fill([x], x)
   return (filled, string_ops.as_string(filled), {
       'structure': string_ops.as_string(filled)
   })
Ejemplo n.º 47
0
 def testFillString(self):
     np_ans = np.array([[b"yolo"] * 3] * 2)
     with self.test_session(use_gpu=False):
         tf_ans = array_ops.fill([2, 3], np_ans[0][0], name="fill").eval()
     self.assertAllEqual(np_ans, tf_ans)
Ejemplo n.º 48
0
def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size,
                      beam_width, end_token, length_penalty_weight,
                      coverage_penalty_weight):
    """Performs a single step of Beam Search Decoding.
  Args:
    time: Beam search time step, should start at 0. At time 0 we assume
      that all beams are equal and consider only the first beam for
      continuations.
    logits: Logits at the current time step. A tensor of shape
      `[batch_size, beam_width, vocab_size]`
    next_cell_state: The next state from the cell, e.g. an instance of
      AttentionWrapperState if the cell is attentional.
    beam_state: Current state of the beam search.
      An instance of `BeamSearchDecoderState`.
    batch_size: The batch size for this input.
    beam_width: Python int.  The size of the beams.
    end_token: The int32 end token.
    length_penalty_weight: Float weight to penalize length. Disabled with 0.0.
    coverage_penalty_weight: Float weight to penalize the coverage of source
      sentence. Disabled with 0.0.
  Returns:
    A new beam state.
  """
    static_batch_size = tensor_util.constant_value(batch_size)

    # Calculate the current lengths of the predictions
    prediction_lengths = beam_state.lengths
    previously_finished = beam_state.finished
    not_finished = math_ops.logical_not(previously_finished)

    # Calculate the total log probs for the new hypotheses
    # Final Shape: [batch_size, beam_width, vocab_size]
    step_log_probs = nn_ops.log_softmax(logits)
    step_log_probs = _mask_probs(step_log_probs, end_token,
                                 previously_finished)
    total_probs = array_ops.expand_dims(beam_state.log_probs,
                                        2) + step_log_probs

    # Calculate the continuation lengths by adding to all continuing beams.
    vocab_size = logits.shape[-1].value or array_ops.shape(logits)[-1]
    lengths_to_add = array_ops.one_hot(indices=array_ops.fill(
        [batch_size, beam_width], end_token),
                                       depth=vocab_size,
                                       on_value=np.int64(0),
                                       off_value=np.int64(1),
                                       dtype=dtypes.int64)
    add_mask = math_ops.to_int64(not_finished)
    lengths_to_add *= array_ops.expand_dims(add_mask, 2)
    new_prediction_lengths = (lengths_to_add +
                              array_ops.expand_dims(prediction_lengths, 2))

    # Calculate the accumulated attention probabilities if coverage penalty is
    # enabled.
    accumulated_attention_probs = None
    attention_probs = get_attention_probs(next_cell_state,
                                          coverage_penalty_weight)
    if attention_probs is not None:
        attention_probs *= array_ops.expand_dims(
            math_ops.to_float(not_finished), 2)
        accumulated_attention_probs = (beam_state.accumulated_attention_probs +
                                       attention_probs)

    # Calculate the scores for each beam
    scores = _get_scores(
        log_probs=total_probs,
        sequence_lengths=new_prediction_lengths,
        length_penalty_weight=length_penalty_weight,
        coverage_penalty_weight=coverage_penalty_weight,
        finished=previously_finished,
        accumulated_attention_probs=accumulated_attention_probs)

    time = ops.convert_to_tensor(time, name="time")
    # During the first time step we only consider the initial beam
    scores_flat = array_ops.reshape(scores, [batch_size, -1])

    # Pick the next beams according to the specified successors function
    next_beam_size = ops.convert_to_tensor(beam_width,
                                           dtype=dtypes.int32,
                                           name="beam_width")
    next_beam_scores, word_indices = nn_ops.top_k(scores_flat,
                                                  k=next_beam_size)

    next_beam_scores.set_shape([static_batch_size, beam_width])
    word_indices.set_shape([static_batch_size, beam_width])

    # Pick out the probs, beam_ids, and states according to the chosen predictions
    next_beam_probs = _tensor_gather_helper(gather_indices=word_indices,
                                            gather_from=total_probs,
                                            batch_size=batch_size,
                                            range_size=beam_width * vocab_size,
                                            gather_shape=[-1],
                                            name="next_beam_probs")
    # Note: just doing the following
    #   math_ops.to_int32(word_indices % vocab_size,
    #       name="next_beam_word_ids")
    # would be a lot cleaner but for reasons unclear, that hides the results of
    # the op which prevents capturing it with tfdbg debug ops.
    raw_next_word_ids = math_ops.mod(word_indices,
                                     vocab_size,
                                     name="next_beam_word_ids")
    next_word_ids = math_ops.to_int32(raw_next_word_ids)
    next_beam_ids = math_ops.to_int32(word_indices / vocab_size,
                                      name="next_beam_parent_ids")

    # Append new ids to current predictions
    previously_finished = _tensor_gather_helper(
        gather_indices=next_beam_ids,
        gather_from=previously_finished,
        batch_size=batch_size,
        range_size=beam_width,
        gather_shape=[-1])
    next_finished = math_ops.logical_or(previously_finished,
                                        math_ops.equal(next_word_ids,
                                                       end_token),
                                        name="next_beam_finished")

    # Calculate the length of the next predictions.
    # 1. Finished beams remain unchanged.
    # 2. Beams that are now finished (EOS predicted) have their length
    #    increased by 1.
    # 3. Beams that are not yet finished have their length increased by 1.
    lengths_to_add = math_ops.to_int64(
        math_ops.logical_not(previously_finished))
    next_prediction_len = _tensor_gather_helper(gather_indices=next_beam_ids,
                                                gather_from=beam_state.lengths,
                                                batch_size=batch_size,
                                                range_size=beam_width,
                                                gather_shape=[-1])
    next_prediction_len += lengths_to_add
    next_accumulated_attention_probs = ()
    if accumulated_attention_probs is not None:
        next_accumulated_attention_probs = _tensor_gather_helper(
            gather_indices=next_beam_ids,
            gather_from=accumulated_attention_probs,
            batch_size=batch_size,
            range_size=beam_width,
            gather_shape=[batch_size * beam_width, -1],
            name="next_accumulated_attention_probs")

    # Pick out the cell_states according to the next_beam_ids. We use a
    # different gather_shape here because the cell_state tensors, i.e.
    # the tensors that would be gathered from, all have dimension
    # greater than two and we need to preserve those dimensions.
    # pylint: disable=g-long-lambda
    next_cell_state = nest.map_structure(
        lambda gather_from: _maybe_tensor_gather_helper(
            gather_indices=next_beam_ids,
            gather_from=gather_from,
            batch_size=batch_size,
            range_size=beam_width,
            gather_shape=[batch_size * beam_width, -1]), next_cell_state)
    # pylint: enable=g-long-lambda
    next_state = BeamSearchDecoderState(
        cell_state=next_cell_state,
        log_probs=next_beam_probs,
        lengths=next_prediction_len,
        finished=next_finished,
        accumulated_attention_probs=next_accumulated_attention_probs,
    )

    output = BeamSearchDecoderOutput(scores=next_beam_scores,
                                     predicted_ids=next_word_ids,
                                     logits=logits,
                                     parent_ids=next_beam_ids)

    return output, next_state
Ejemplo n.º 49
0
 def if_false():
     return array_ops.fill([p], 5.)
Ejemplo n.º 50
0
 def _sparse(i):
     return sparse_tensor.SparseTensorValue(
         indices=array_ops.expand_dims(
             math_ops.range(i, dtype=dtypes.int64), 1),
         values=array_ops.fill([math_ops.cast(i, dtypes.int32)], i),
         dense_shape=[i])
def _embedding_lookup_with_distributed_aggregation(params,
                                                   ids,
                                                   partition_strategy="mod",
                                                   name=None,
                                                   max_norm=None,
                                                   weights=None,
                                                   idx=None,
                                                   segment_ids=None):
    """Lookup helper for embedding_lookup_sparse_with_distributed_aggregation."""
    if params is None or params == []:  # pylint: disable=g-explicit-bool-comparison
        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]

    def maybe_normalize(x):
        if max_norm is not None:
            if x.get_shape().ndims is not None:
                ndims = x.get_shape().ndims
            else:
                ndims = array_ops.size(array_ops.shape(x))
            return clip_ops.clip_by_norm(x,
                                         max_norm,
                                         axes=list(range(1, ndims)))
        return x

    with ops.name_scope(name, "embedding_lookup_with_distributed_aggregation",
                        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")
        if np == 1:
            with ops.colocate_with(params[0]):
                ret = maybe_normalize(_do_gather(params[0], ids))
                ignore_weights = weights is None
                if not ignore_weights:
                    if weights.dtype != ret.dtype:
                        weights = math_ops.cast(weights, ret.dtype)
                    # Reshape to allow broadcast
                    ones = array_ops.fill(
                        array_ops.expand_dims(array_ops.rank(ret) - 1, 0), 1)
                    bcast_weights_shape = array_ops.concat(
                        [array_ops.shape(weights), ones], 0)
                    orig_weights_shape = weights.get_shape()
                    weights = array_ops.reshape(weights, bcast_weights_shape)
                    # Set weights shape after reshape
                    if ret.get_shape().ndims is not None:
                        weights.set_shape(
                            orig_weights_shape.concatenate(
                                [1 for _ in range(ret.get_shape().ndims - 1)]))
                    ret *= weights
                    return math_ops.segment_sum(ret, segment_ids, name=name)
                else:
                    return math_ops.sparse_segment_sum(ret,
                                                       idx,
                                                       segment_ids,
                                                       name=name)
        else:
            ids = ops.convert_to_tensor(ids, name="ids")
            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 = params[0].get_shape().dims[0]
                for p in xrange(1, np):
                    dim_0_size += params[p].get_shape().dims[0]
                if dim_0_size.value:
                    num_total_ids = constant_op.constant(
                        dim_0_size, flat_ids.dtype)
                else:
                    dim_0_sizes = []
                    for p in xrange(np):
                        if params[p].get_shape().dims[0].value is not None:
                            dim_0_sizes.append(
                                params[p].get_shape().dims[0].value)
                        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
                is_in_first_extras_partitions = math_ops.cast(
                    p_assignments < extras, flat_ids.dtype)
                new_ids = (is_in_first_extras_partitions *
                           (flat_ids % (ids_per_partition + 1)) +
                           (1 - is_in_first_extras_partitions) *
                           ((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):
                with ops.colocate_with(params[p]):
                    partitioned_result.append(
                        _do_gather(params[p], gather_ids[p]))

            ignore_weights = weights is None
            if not ignore_weights:
                # Partition weights according to pindices.
                partitioned_weight = []
                for p in xrange(np):
                    partitioned_weight.append(
                        array_ops.gather(weights, pindices[p]))
            # Reshape each partition result.
            element_shape = params[0].get_shape()[1:]
            for p in params[1:]:
                element_shape = element_shape.merge_with(p.get_shape()[1:])
            if element_shape.is_fully_defined():
                for p in xrange(np):
                    with ops.colocate_with(params[p]):
                        partitioned_result[p] = array_ops.reshape(
                            partitioned_result[p],
                            array_ops.concat(
                                [array_ops.shape(pindices[p]), element_shape],
                                0))
            else:
                with ops.colocate_with(params[0]):
                    params_shape = array_ops.shape(params[0])
                for p in xrange(np):
                    with ops.colocate_with(params[p]):
                        partitioned_result[p] = array_ops.reshape(
                            partitioned_result[p],
                            array_ops.concat([
                                array_ops.shape(pindices[p]),
                                array_ops.slice(params_shape, [1], [-1])
                            ], 0))
            # Normalize each partition result.
            for p in xrange(np):
                with ops.colocate_with(params[p]):
                    partitioned_result[p] = maybe_normalize(
                        partitioned_result[p])
            if not ignore_weights:
                # Multiply each partition result with partition weights.
                for p in xrange(np):
                    with ops.colocate_with(params[p]):
                        if partitioned_weight[p].dtype != partitioned_result[
                                p].dtype:
                            partitioned_weight[p] = math_ops.cast(
                                partitioned_weight[p],
                                partitioned_result[p].dtype)
                        # Reshape partition weights.
                        ones = array_ops.fill(
                            array_ops.expand_dims(
                                array_ops.rank(partitioned_result[p]) - 1, 0),
                            1)
                        bcast_weights_shape = array_ops.concat(
                            [array_ops.shape(partitioned_weight[p]), ones], 0)
                        orig_weights_shape = partitioned_weight[p].get_shape()
                        partitioned_weight[p] = array_ops.reshape(
                            partitioned_weight[p], bcast_weights_shape)
                        if partitioned_result[p].get_shape().ndims is not None:
                            partitioned_weight[p].set_shape(
                                orig_weights_shape.concatenate([
                                    1 for _ in range(partitioned_result[p].
                                                     get_shape().ndims - 1)
                                ]))
                        partitioned_result[p] *= partitioned_weight[p]
            partitioned_segment_ids = []
            for p in xrange(np):
                if not ignore_weights:
                    # Partition segment_ids according to pindices.
                    p_segment_ids = array_ops.gather(segment_ids, pindices[p])
                    # Number the p_segment_ids to meet segment_sum's requirements. Note
                    # that unique_p_segment_ids contains unique segment ids of this
                    # partition and these ids' order is unchanged.
                    unique_p_segment_ids, unique_p_segment_idx = array_ops.unique(
                        p_segment_ids)
                    partitioned_segment_ids.append(unique_p_segment_ids)
                    # segment_sum this partition's result.
                    with ops.colocate_with(params[p]):
                        partitioned_result[p] = math_ops.segment_sum(
                            partitioned_result[p], unique_p_segment_idx)
                else:
                    # When ignore weights, we need to get indexs of elements in idx and
                    # segment_ids.
                    _, exclude_idx = array_ops.setdiff1d(idx, pindices[p])
                    all_idx = math_ops.range(array_ops.shape(idx)[0])
                    _, include_idx = array_ops.setdiff1d(all_idx, exclude_idx)
                    # Gather segment_ids and idx according to indexs.
                    p_segment_ids = array_ops.gather(segment_ids, include_idx)
                    p_idx = array_ops.gather(idx, include_idx)
                    # Number the p_segment_ids, same as ignore_weights case above.
                    unique_p_segment_ids, unique_p_segment_idx = array_ops.unique(
                        p_segment_ids)
                    _, unique_p_idx_idx = array_ops.unique(p_idx)
                    partitioned_segment_ids.append(unique_p_segment_ids)
                    with ops.colocate_with(params[p]):
                        partitioned_result[p] = math_ops.sparse_segment_sum(
                            partitioned_result[p], unique_p_idx_idx,
                            unique_p_segment_idx)
            # Concat each partition's segment_ids and result for final segment_sum.
            concat_segment_ids = array_ops.concat(partitioned_segment_ids, 0)
            concat_partitioned_result = array_ops.concat(partitioned_result, 0)
            return math_ops.unsorted_segment_sum(
                concat_partitioned_result,
                concat_segment_ids,
                math_ops.reduce_max(concat_segment_ids) + 1,
                name=name)
Ejemplo n.º 52
0
def sparse_to_indicator(sp_input, vocab_size, name=None):
  """Converts a `SparseTensor` of ids into a dense bool indicator tensor.

  The last dimension of `sp_input` is discarded and replaced with the values of
  `sp_input`.  If `sp_input.shape = [D0, D1, ..., Dn, K]`, then
  `output.shape = [D0, D1, ..., Dn, vocab_size]`, where

      output[d_0, d_1, ..., d_n, sp_input[d_0, d_1, ..., d_n, k]] = True

  and False elsewhere in `output`.

  For example, if `sp_input.shape = [2, 3, 4]` with non-empty values:

      [0, 0, 0]: 0
      [0, 1, 0]: 10
      [1, 0, 3]: 103
      [1, 1, 2]: 150
      [1, 1, 3]: 149
      [1, 1, 4]: 150
      [1, 2, 1]: 121

  and `vocab_size = 200`, then the output will be a `[2, 3, 200]` dense bool
  tensor with False everywhere except at positions

      (0, 0, 0), (0, 1, 10), (1, 0, 103), (1, 1, 149), (1, 1, 150),
      (1, 2, 121).

  Note that repeats are allowed in the input SparseTensor.
  This op is useful for converting `SparseTensor`s into dense formats for
  compatibility with ops that expect dense tensors.

  The input `SparseTensor` must be in row-major order.

  Args:
    sp_input: A `SparseTensor` of type `int32` or `int64`.
    vocab_size: The new size of the last dimension, with
      `all(0 <= sp_input.values < vocab_size)`.
    name: A name prefix for the returned tensors (optional)

  Returns:
    A dense bool indicator tensor representing the indices with specified value.

  Raises:
    TypeError: If `sp_input` is not a `SparseTensor`.
  """
  if not isinstance(sp_input, ops.SparseTensor):
    raise TypeError("Input must be a SparseTensor")

  with ops.op_scope([sp_input], name, "SparseToIndicator") as name:
    indices_shape = array_ops.shape(sp_input.indices)
    num_entries = indices_shape[0]
    rank = indices_shape[1]

    ids = sp_input.values
    if ids.dtype != dtypes.int64:
      ids = math_ops.cast(ids, dtypes.int64)

    # Slice off the last dimension of indices, then then tack on the ids
    indices_columns_to_preserve = array_ops.slice(
        sp_input.indices, [0, 0], array_ops.pack([-1, rank - 1]))
    new_indices = array_ops.concat(1, [indices_columns_to_preserve,
                                       array_ops.reshape(ids, [-1, 1])])

    new_values = array_ops.fill(array_ops.expand_dims(num_entries, 0), True)
    new_shape = array_ops.concat(0, [array_ops.slice(
        sp_input.shape, [0], array_ops.expand_dims(rank - 1, 0)), [vocab_size]])

    sp_new = ops.SparseTensor(new_indices, new_values, new_shape)

    # validate_indices may be False because we allow duplicates in new_indices:
    # repeated indices are allowed when creating an indicator matrix.
    return sparse_tensor_to_dense(
        sp_new, default_value=False, validate_indices=False, name=name)
Ejemplo n.º 53
0
 def _compare(self, dims, val, np_ans, use_gpu):
     with self.test_session(use_gpu=use_gpu):
         tf_ans = array_ops.fill(dims, val, name="fill")
         out = tf_ans.eval()
     self.assertAllClose(np_ans, out)
Ejemplo n.º 54
0
 def f(x):
     if x.dtype == dtypes.bool:
         return array_ops.fill(array_ops.shape(x), False)
     return x < 0
Ejemplo n.º 55
0
 def body(output, i):
     value = array_ops.fill(counts[i:i + 1], i)
     return (output.write(i, value), i + 1)
def embedding_lookup_sparse_with_distributed_aggregation(
        params,
        sp_ids,
        sp_weights,
        partition_strategy="mod",
        name=None,
        combiner=None,
        max_norm=None):
    """Computes embeddings for the given ids and weights.

  Embeddings belonging to same param are aggregated on that device first. This
  op is intended to decrease data transmission and improve parallelism. See
  `tf.nn.embedding_lookup_sparse` for the functionality and example of this op.

  Args:
    params: A single tensor representing the complete embedding tensor, or a
      list of P tensors all of same shape except for the first dimension,
      representing sharded embedding tensors.  Alternatively, a
      `PartitionedVariable`, created by partitioning along dimension 0. Each
      element must be appropriately sized for the given `partition_strategy`.
    sp_ids: N x M SparseTensor of int64 ids (typically from FeatureValueToId),
      where N is typically batch size and M is arbitrary.
    sp_weights: either a SparseTensor of float / double weights, or None to
      indicate all weights should be taken to be 1. If specified, sp_weights
      must have exactly the same shape and indices as sp_ids.
    partition_strategy: A string specifying the partitioning strategy, relevant
      if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
      is `"mod"`. See `tf.nn.embedding_lookup` for more details.
    name: Optional name for the op.
    combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
      and "sum" are supported. "sum" computes the weighted sum of the embedding
      results for each row. "mean" is the weighted sum divided by the total
      weight. "sqrtn" is the weighted sum divided by the square root of the sum
      of the squares of the weights.
    max_norm: If not None, each embedding is normalized to have l2 norm equal to
      max_norm before combining.

  Returns:
    A dense tensor representing the combined embeddings for the
    sparse ids. For each row in the dense tensor represented by sp_ids, the op
    looks up the embeddings for all ids in that row, multiplies them by the
    corresponding weight, and combines these embeddings as specified.

  Raises:
    TypeError: If sp_ids is not a SparseTensor, or if sp_weights is neither
      None nor SparseTensor.
    ValueError: If combiner is not one of {"mean", "sqrtn", "sum"}.
  """
    if combiner is None:
        logging.warn("The default value of combiner will change from \"mean\" "
                     "to \"sqrtn\" after 2016/11/01.")
        combiner = "mean"
    if combiner not in ("mean", "sqrtn", "sum"):
        raise ValueError("combiner must be one of 'mean', 'sqrtn' or 'sum'")
    if isinstance(params, variables.PartitionedVariable):
        params = list(params)  # Iterate to get the underlying Variables.
    if not isinstance(params, list):
        params = [params]
    if not isinstance(sp_ids, sparse_tensor.SparseTensor):
        raise TypeError("sp_ids must be SparseTensor")
    ignore_weights = sp_weights is None
    if not ignore_weights:
        if not isinstance(sp_weights, sparse_tensor.SparseTensor):
            raise TypeError("sp_weights must be either None or SparseTensor")
        sp_ids.values.get_shape().assert_is_compatible_with(
            sp_weights.values.get_shape())
        sp_ids.indices.get_shape().assert_is_compatible_with(
            sp_weights.indices.get_shape())
        sp_ids.dense_shape.get_shape().assert_is_compatible_with(
            sp_weights.dense_shape.get_shape())
        # TODO(yleon): Add enhanced node assertions to verify that sp_ids and
        # sp_weights have equal indices and shapes.

    with ops.name_scope(name, "embedding_lookup_sparse",
                        params + [sp_ids]) as name:
        segment_ids = sp_ids.indices[:, 0]
        if segment_ids.dtype != dtypes.int32:
            segment_ids = math_ops.cast(segment_ids, dtypes.int32)

        ids = sp_ids.values
        if ignore_weights:
            ids, idx = array_ops.unique(ids)
        else:
            idx = None

        weights = None if ignore_weights else sp_weights.values
        embeddings = _embedding_lookup_with_distributed_aggregation(
            params,
            ids,
            partition_strategy=partition_strategy,
            max_norm=max_norm,
            weights=weights,
            idx=idx,
            segment_ids=segment_ids)
        # Set weights to all one if ignore weights.
        if ignore_weights:
            weights = array_ops.fill([array_ops.shape(segment_ids)[0]], 1)
        if weights.dtype != embeddings.dtype:
            weights = math_ops.cast(weights, embeddings.dtype)
        # Reshape weights.
        ones = array_ops.fill(
            array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
        bcast_weights_shape = array_ops.concat(
            [array_ops.shape(weights), ones], 0)
        orig_weights_shape = weights.get_shape()
        weights = array_ops.reshape(weights, bcast_weights_shape)
        if embeddings.get_shape().ndims is not None:
            weights.set_shape(
                orig_weights_shape.concatenate(
                    [1 for _ in range(embeddings.get_shape().ndims - 1)]))

        if combiner == "mean":
            weight_sum = math_ops.segment_sum(weights, segment_ids)
            embeddings = math_ops.div(embeddings, weight_sum)
        elif combiner == "sqrtn":
            weights_squared = math_ops.pow(weights, 2)
            weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
            weight_sum_sqrt = math_ops.sqrt(weight_sum)
            embeddings = math_ops.div(embeddings, weight_sum_sqrt)
        elif combiner != "sum":
            assert False, "Unrecognized combiner"
        return embeddings
Ejemplo n.º 57
0
 def loop_fn(i):
     return array_ops.fill((2, 3), i)
Ejemplo n.º 58
0
 def _stddev(self):
   if self.allow_nan_stats:
     return array_ops.fill(self.batch_shape_tensor(),
                           self.dtype.as_numpy_dtype(np.nan))
   else:
     raise ValueError("`stddev` is undefined for Cauchy distribution.")
Ejemplo n.º 59
0
    def _testDynamicDecodeRNN(self, time_major, has_attention):
        encoder_sequence_length = [3, 2, 3, 1, 1]
        decoder_sequence_length = [2, 0, 1, 2, 3]
        batch_size = 5
        decoder_max_time = 4
        input_depth = 7
        cell_depth = 9
        attention_depth = 6
        vocab_size = 20
        end_token = vocab_size - 1
        start_token = 0
        embedding_dim = 50
        max_out = max(decoder_sequence_length)
        output_layer = layers_core.Dense(vocab_size,
                                         use_bias=True,
                                         activation=None)
        beam_width = 3

        with self.test_session() as sess:
            batch_size_tensor = constant_op.constant(batch_size)
            embedding = np.random.randn(vocab_size,
                                        embedding_dim).astype(np.float32)
            cell = rnn_cell.LSTMCell(cell_depth)
            if has_attention:
                inputs = array_ops.placeholder_with_default(
                    np.random.randn(batch_size, decoder_max_time,
                                    input_depth).astype(np.float32),
                    shape=(None, None, input_depth))
                tiled_inputs = beam_search_decoder.tile_batch(
                    inputs, multiplier=beam_width)
                tiled_sequence_length = beam_search_decoder.tile_batch(
                    encoder_sequence_length, multiplier=beam_width)
                attention_mechanism = attention_wrapper.BahdanauAttention(
                    num_units=attention_depth,
                    memory=tiled_inputs,
                    memory_sequence_length=tiled_sequence_length)
                cell = attention_wrapper.AttentionWrapper(
                    cell=cell,
                    attention_mechanism=attention_mechanism,
                    attention_layer_size=attention_depth,
                    alignment_history=False)
            cell_state = cell.zero_state(dtype=dtypes.float32,
                                         batch_size=batch_size_tensor *
                                         beam_width)
            bsd = beam_search_decoder.BeamSearchDecoder(
                cell=cell,
                embedding=embedding,
                start_tokens=array_ops.fill([batch_size_tensor], start_token),
                end_token=end_token,
                initial_state=cell_state,
                beam_width=beam_width,
                output_layer=output_layer,
                length_penalty_weight=0.0)

            final_outputs, final_state, final_sequence_lengths = (
                decoder.dynamic_decode(bsd,
                                       output_time_major=time_major,
                                       maximum_iterations=max_out))

            def _t(shape):
                if time_major:
                    return (shape[1], shape[0]) + shape[2:]
                return shape

            self.assertTrue(
                isinstance(final_outputs,
                           beam_search_decoder.FinalBeamSearchDecoderOutput))
            self.assertTrue(
                isinstance(final_state,
                           beam_search_decoder.BeamSearchDecoderState))

            beam_search_decoder_output = final_outputs.beam_search_decoder_output
            self.assertEqual(
                _t((batch_size, None, beam_width)),
                tuple(beam_search_decoder_output.scores.get_shape().as_list()))
            self.assertEqual(
                _t((batch_size, None, beam_width)),
                tuple(final_outputs.predicted_ids.get_shape().as_list()))

            sess.run(variables.global_variables_initializer())
            sess_results = sess.run({
                'final_outputs':
                final_outputs,
                'final_state':
                final_state,
                'final_sequence_lengths':
                final_sequence_lengths
            })

            max_sequence_length = np.max(
                sess_results['final_sequence_lengths'])

            # A smoke test
            self.assertEqual(
                _t((batch_size, max_sequence_length, beam_width)),
                sess_results['final_outputs'].beam_search_decoder_output.
                scores.shape)
            self.assertEqual(
                _t((batch_size, max_sequence_length, beam_width)),
                sess_results['final_outputs'].beam_search_decoder_output.
                predicted_ids.shape)
Ejemplo n.º 60
0
  def _apply_fn(dataset):
    """Function from `Dataset` to `Dataset` that applies the transformation."""
    dist_estimation_batch_size = 32
    target_dist_t = ops.convert_to_tensor(target_dist, name="initial_dist")
    class_values_ds = dataset.map(class_func)
    if initial_dist is not None:
      initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
      acceptance_dist = _calculate_acceptance_probs(initial_dist_t,
                                                    target_dist_t)
      initial_dist_ds = dataset_ops.Dataset.from_tensors(
          initial_dist_t).repeat()
      acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
          acceptance_dist).repeat()
    else:
      num_classes = (target_dist_t.shape[0].value or
                     array_ops.shape(target_dist_t)[0])
      smoothing_constant = 10
      # Disable device functions and colocation constraints so that the variable
      # will be placed with the eventual DT_VARIANT dataset tensor.
      with ops.colocate_with(None, ignore_existing=True):
        num_examples_per_class_seen = resource_variable_ops.ResourceVariable(
            initial_value=array_ops.fill([num_classes],
                                         np.int64(smoothing_constant)),
            trainable=False,
            collections=[ops.GraphKeys.LOCAL_VARIABLES],
            name="local_class_count",
            dtype=dtypes.int64)

      def update_estimate_and_tile(c):
        return array_ops.tile(
            array_ops.expand_dims(
                _estimate_data_distribution(c, num_examples_per_class_seen), 0),
            [dist_estimation_batch_size, 1])

      initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size)
                         .map(update_estimate_and_tile).apply(batching
                                                              .unbatch()))
      acceptance_dist_ds = initial_dist_ds.map(
          lambda initial: _calculate_acceptance_probs(initial, target_dist_t))

    def maybe_warn_on_large_rejection(accept_dist, initial_dist):
      proportion_rejected = math_ops.reduce_sum(
          (1 - accept_dist) * initial_dist)
      return control_flow_ops.cond(
          math_ops.less(proportion_rejected, .5),
          lambda: accept_dist,
          lambda: logging_ops.Print(  # pylint: disable=g-long-lambda
              accept_dist, [proportion_rejected, initial_dist, accept_dist],
              message="Proportion of examples rejected by sampler is high: ",
              summarize=100,
              first_n=10))

    acceptance_dist_ds = (dataset_ops.Dataset.zip((acceptance_dist_ds,
                                                   initial_dist_ds))
                          .map(maybe_warn_on_large_rejection))

    current_probabilities_ds = dataset_ops.Dataset.zip(
        (acceptance_dist_ds, class_values_ds)).map(array_ops.gather)
    filtered_ds = (
        dataset_ops.Dataset.zip((class_values_ds, current_probabilities_ds,
                                 dataset))
        .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))
    return filtered_ds.map(lambda class_value, _, data: (class_value, data))