Python math_ops.segment_sum函数代码示例

    def body(it, cost):
      embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
      cost = control_flow_ops.cond(
          math_ops.equal(it, 3), lambda: math_ops.square(cost),
          (lambda: cost + math_ops.reduce_sum(embedding)))
      return it + 1, cost

      _, cost = control_flow_ops.while_loop(
          cond, body, [constant_op.constant(0),
                       constant_op.constant(0.0)])

      dynamic_grads = gradients_impl.gradients(cost, [embedding_matrix])[0]
      dynamic_grads = math_ops.segment_sum(dynamic_grads.values,
                                           dynamic_grads.indices)

      embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
      static = math_ops.square(
          math_ops.reduce_sum(embedding) + math_ops.reduce_sum(embedding) +
          math_ops.reduce_sum(embedding)) + math_ops.reduce_sum(embedding)
      static_grads = gradients_impl.gradients(static, [embedding_matrix])[0]
      static_grads = math_ops.segment_sum(static_grads.values,
                                          static_grads.indices)

      with self.cached_session():
        self.evaluate(variables.global_variables_initializer())
        self.assertAllEqual(*self.evaluate([static_grads, dynamic_grads]))

  def doTestIndexedSlicesGradientInCondInWhileLoop(self, use_resource=False):
    with ops.Graph().as_default():
      embedding_matrix = variable_scope.get_variable(
          "embedding_matrix", [5, 5],
          initializer=init_ops.random_normal_initializer(),
          use_resource=use_resource)

      def Cond(it, _):
        return it < 5

      def Body(it, cost):
        embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
        cost = control_flow_ops.cond(
            math_ops.equal(it, 3), lambda: math_ops.square(cost),
            lambda: cost + math_ops.reduce_sum(embedding))
        return it + 1, cost

      _, cost = control_flow_ops.while_loop(
          Cond, Body, [constant_op.constant(0), constant_op.constant(0.0)])

      dynamic_grads = gradients_impl.gradients(cost, [embedding_matrix])[0]
      dynamic_grads = math_ops.segment_sum(dynamic_grads.values,
                                           dynamic_grads.indices)

      embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
      static = math_ops.square(
          math_ops.reduce_sum(embedding) + math_ops.reduce_sum(embedding) +
          math_ops.reduce_sum(embedding)) + math_ops.reduce_sum(embedding)
      static_grads = gradients_impl.gradients(static, [embedding_matrix])[0]
      static_grads = math_ops.segment_sum(static_grads.values,
                                          static_grads.indices)

      with self.test_session() as sess:
        sess.run(variables.global_variables_initializer())
        self.assertAllEqual(*sess.run([static_grads, dynamic_grads]))

  def testGradientMatchesSegmentSum(self):
    # Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum
    # and compare the outputs, which should be identical.
    # NB: for this test to work, indices must be valid for SegmentSum, namely
    # it must be sorted, the indices must be contiguous, and num_segments
    # must be max(indices) + 1.
    indices = [0, 0, 1, 1, 1, 2, 3, 4, 5]
    n = len(indices)
    num_cols = 2
    shape = [n, num_cols]
    num_segments = max(indices) + 1
    for dtype in self.differentiable_dtypes:
      with self.cached_session(use_gpu=True):
        tf_x, np_x = self._input(shape, dtype=dtype)
        # Results from UnsortedSegmentSum
        unsorted_s = math_ops.unsorted_segment_sum(
            data=tf_x, segment_ids=indices, num_segments=num_segments)
        unsorted_jacob_t, unsorted_jacob_n = (
            gradient_checker.compute_gradient(tf_x, shape, unsorted_s,
                                              [num_segments, num_cols],
                                              x_init_value=np_x, delta=1))

        # Results from SegmentSum
        sorted_s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
        sorted_jacob_t, sorted_jacob_n = gradient_checker.compute_gradient(
            tf_x,
            shape,
            sorted_s, [num_segments, num_cols],
            x_init_value=np_x,
            delta=1)
      self.assertAllClose(unsorted_jacob_t, sorted_jacob_t)
      self.assertAllClose(unsorted_jacob_n, sorted_jacob_n)

 def testSegmentIdsInvalid2(self):
   shape = [4, 4]
   with self.test_session():
     tf_x, _ = self._input(shape)
     indices = [1, 1, 2, 2]
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     with self.assertRaisesOpError("segment ids do not start at 0"):
       s.eval()

 def testSegmentIdsInvalid2(self):
   shape = [4, 4]
   with self.cached_session():
     tf_x, _ = self._input(shape)
     indices = [0, 1, 0, 1]
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     with self.assertRaisesOpError("segment ids are not increasing"):
       s.eval()

 def testSegmentIdsInvalid5(self):
   shape = [4, 4]
   with self.test_session():
     tf_x, _ = self._input(shape)
     indices = [0, 0, 0, -2]
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     with self.assertRaisesOpError("segment ids must be >= 0"):
       s.eval()

 def testSegmentIdsValid(self):
   # This is a baseline for the following SegmentIdsInvalid* tests.
   shape = [4, 4]
   with self.test_session():
     tf_x, _ = self._input(shape)
     indices = [0, 0, 0, 1]
     result = math_ops.segment_sum(data=tf_x, segment_ids=indices).eval()
     self.assertAllEqual([[15, 18, 21, 24], [13, 14, 15, 16]], result)

 def testSegmentIdsSize(self):
   shape = [4, 4]
   with self.test_session():
     tf_x, _ = self._input(shape)
     indices = [0, 1]
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     with self.assertRaisesOpError("segment_ids should be the same size"):
       s.eval()

 def testSegmentIdsInvalid5(self):
   shape = [4, 4]
   for use_gpu in [True, False]:
     with self.cached_session(use_gpu=use_gpu):
       tf_x, _ = self._input(shape, dtype=dtypes_lib.float32)
       indices = [0, 0, 0, -2]
       s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
       with self.assertRaisesOpError("segment ids must be >= 0"):
         s.eval()

 def testSegmentIdsValid(self):
   # This is a baseline for the following SegmentIdsInvalid* tests.
   shape = [4, 4]
   for use_gpu in [True, False]:
     with self.cached_session(use_gpu=use_gpu):
       tf_x, _ = self._input(shape, dtype=dtypes_lib.float32)
       indices = [0, 0, 0, 1]
       result = math_ops.segment_sum(data=tf_x, segment_ids=indices).eval()
       self.assertAllEqual([[15, 18, 21, 24], [13, 14, 15, 16]], result)

 def testSegmentIdsSize(self):
   shape = [4, 4]
   for use_gpu in [True, False]:
     with self.cached_session(use_gpu=use_gpu):
       tf_x, _ = self._input(shape)
       indices = [0, 1]
       s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
       with self.assertRaisesOpError("segment_ids should be the same size"):
         s.eval()

 def testSegmentIdsHole(self):
   shape = [4, 4]
   with self.test_session():
     tf_x, np_x = self._input(shape)
     indices = [0, 0, 3, 3]
     np_ans = self._segmentReduce(indices, np_x, np.add)
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     tf_ans = s.eval()
     self.assertAllClose(np_ans, tf_ans)

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

 def testSegmentIdsHole(self):
   shape = [4, 4]
   for use_gpu in [True, False]:
     with self.cached_session(use_gpu=use_gpu):
       tf_x, np_x = self._input(shape, dtype=dtypes_lib.float32)
       indices = [0, 0, 3, 3]
       np_ans = self._segmentReduce(indices, np_x, np.add)
       s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
       tf_ans = s.eval()
       self.assertAllClose(np_ans, tf_ans)

 def testSegmentIdsInvalid3(self):
   shape = [4, 4]
   with self.cached_session():
     tf_x, _ = self._input(shape)
     indices = [0, 1, 2, 0]
     s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
     with self.assertRaisesOpError(
         r"Segment id 1 out of range \[0, 1\), possibly "
         "because 'segment_ids' input is not sorted."):
       s.eval()