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()

def _SegmentMinOrMaxGrad(op, grad):
  """ Gradient for SegmentMin and SegmentMax. """
  zeros = array_ops.zeros_like(op.inputs[0], dtype=op.inputs[0].dtype)
  # Get the number of selected (minimum or maximum) elements in each segment.
  gathered_outputs = array_ops.gather(op.outputs[0], op.inputs[1])
  is_selected = math_ops.equal(op.inputs[0], gathered_outputs)
  num_selected = math_ops.segment_sum(math_ops.cast(is_selected, grad.dtype),
                                      op.inputs[1])
  # Compute the gradient for each segment. The gradient for the ith segment is
  # divided evenly among the selected elements in that segment.
  weighted_grads = math_ops.div(grad, num_selected)
  gathered_grads = array_ops.gather(weighted_grads, op.inputs[1])
  return array_ops.where(is_selected, gathered_grads, zeros), None

    def _linear_predictions(self, examples):
        """Returns predictions of the form w*x."""
        with name_scope("sdca/prediction"):
            sparse_variables = self._convert_n_to_tensor(self._variables["sparse_features_weights"])
            result = 0.0
            for sfc, sv in zip(examples["sparse_features"], sparse_variables):
                # TODO(sibyl-Aix6ihai): following does not take care of missing features.
                result += math_ops.segment_sum(
                    math_ops.mul(array_ops.gather(sv, sfc.feature_indices), sfc.feature_values), sfc.example_indices
                )
            dense_features = self._convert_n_to_tensor(examples["dense_features"])
            dense_variables = self._convert_n_to_tensor(self._variables["dense_features_weights"])

            for i in range(len(dense_variables)):
                result += math_ops.matmul(dense_features[i], array_ops.expand_dims(dense_variables[i], -1))

        # Reshaping to allow shape inference at graph construction time.
        return array_ops.reshape(result, [-1])

 def _linear_predictions(self, examples):
   """Returns predictions of the form w*x."""
   with name_scope('sdca/prediction'):
     sparse_variables = self._convert_n_to_tensor(self._variables[
         'sparse_features_weights'])
     predictions = 0
     for st_i, sv in zip(examples['sparse_features'], sparse_variables):
       ei, fi = array_ops.split(1, 2, st_i.indices)
       ei = array_ops.reshape(ei, [-1])
       fi = array_ops.reshape(fi, [-1])
       fv = array_ops.reshape(st_i.values, [-1])
       # TODO(rohananil): This does not work if examples have empty features.
       predictions += math_ops.segment_sum(
           math_ops.mul(
               array_ops.gather(sv, fi), fv), array_ops.reshape(ei, [-1]))
     dense_features = self._convert_n_to_tensor(examples['dense_features'])
     dense_variables = self._convert_n_to_tensor(self._variables[
         'dense_features_weights'])
     for i in range(len(dense_variables)):
       predictions += dense_features[i] * dense_variables[i]
   return predictions

 def _logits(self, examples):
   """Compute logits for each example."""
   with name_scope('logits'):
     sparse_variables = self._convert_n_to_tensor(self._variables[
         'sparse_features_weights'])
     logits = 0
     for st_i, sv in zip(examples['sparse_features'], sparse_variables):
       ei, fi = array_ops.split(1, 2, st_i.indices)
       ei = array_ops.reshape(ei, [-1])
       fi = array_ops.reshape(fi, [-1])
       fv = array_ops.reshape(st_i.values, [-1])
       # TODO(rohananil): This does not work if examples have empty features.
       logits += math_ops.segment_sum(
           math_ops.mul(
               array_ops.gather(sv, fi), fv), array_ops.reshape(ei, [-1]))
     dense_features = self._convert_n_to_tensor(examples['dense_features'])
     dense_variables = self._convert_n_to_tensor(self._variables[
         'dense_features_weights'])
     for i in xrange(len(dense_variables)):
       logits += dense_features[i] * dense_variables[i]
     return logits

  def _linear_predictions(self, examples):
    """Returns predictions of the form w*x."""
    with name_scope('sdca/prediction'):
      sparse_variables = self._convert_n_to_tensor(self._variables[
          'sparse_features_weights'])
      result = 0.0
      for st_i, sv in zip(examples['sparse_features'], sparse_variables):
        ei, fi = array_ops.split(1, 2, st_i.indices)
        ei = array_ops.reshape(ei, [-1])
        fi = array_ops.reshape(fi, [-1])
        fv = array_ops.reshape(st_i.values, [-1])
        # TODO(sibyl-Aix6ihai): This does not work if examples have empty features.
        result += math_ops.segment_sum(
            math_ops.mul(array_ops.gather(sv, fi), fv), ei)
      dense_features = self._convert_n_to_tensor(examples['dense_features'])
      dense_variables = self._convert_n_to_tensor(self._variables[
          'dense_features_weights'])

      for i in range(len(dense_variables)):
        result += dense_features[i] * dense_variables[i]

    # Reshaping to allow shape inference at graph construction time.
    return array_ops.reshape(result, [-1])