nodes: 0, module; 1, function_definition; 2, function_name:_rowwise_unsorted_segment_sum; 3, parameters; 4, block; 5, identifier:values; 6, identifier:indices; 7, identifier:n; 8, expression_statement; 9, expression_statement; 10, expression_statement; 11, expression_statement; 12, return_statement; 13, comment:"""UnsortedSegmentSum on each row.
  
    Args:
      values: a `Tensor` with shape `[batch_size, k]`.
      indices: an integer `Tensor` with shape `[batch_size, k]`.
      n: an integer.
    Returns:
      A `Tensor` with the same type as `values` and shape `[batch_size, n]`.
    """; 14, assignment; 15, assignment; 16, assignment; 17, call; 18, pattern_list; 19, call; 20, identifier:indices_flat; 21, binary_operator:tf.reshape(indices, [-1]) + tf.div(tf.range(batch * k), k) * n; 22, identifier:ret_flat; 23, call; 24, attribute; 25, argument_list; 26, identifier:batch; 27, identifier:k; 28, attribute; 29, argument_list; 30, call; 31, binary_operator:tf.div(tf.range(batch * k), k) * n; 32, attribute; 33, argument_list; 34, identifier:tf; 35, identifier:reshape; 36, identifier:ret_flat; 37, list; 38, identifier:tf; 39, identifier:unstack; 40, call; 41, keyword_argument; 42, attribute; 43, argument_list; 44, call; 45, identifier:n; 46, identifier:tf; 47, identifier:unsorted_segment_sum; 48, call; 49, identifier:indices_flat; 50, binary_operator:batch * n; 51, identifier:batch; 52, identifier:n; 53, attribute; 54, argument_list; 55, identifier:num; 56, integer:2; 57, identifier:tf; 58, identifier:reshape; 59, identifier:indices; 60, list; 61, attribute; 62, argument_list; 63, attribute; 64, argument_list; 65, identifier:batch; 66, identifier:n; 67, identifier:tf; 68, identifier:shape; 69, identifier:indices; 70, unary_operator; 71, identifier:tf; 72, identifier:div; 73, call; 74, identifier:k; 75, identifier:tf; 76, identifier:reshape; 77, identifier:values; 78, list; 79, integer:1; 80, attribute; 81, argument_list; 82, unary_operator; 83, identifier:tf; 84, identifier:range; 85, binary_operator:batch * k; 86, integer:1; 87, identifier:batch; 88, identifier:k
code: def _rowwise_unsorted_segment_sum(values, indices, n):
    """UnsortedSegmentSum on each row.
  
    Args:
      values: a `Tensor` with shape `[batch_size, k]`.
      indices: an integer `Tensor` with shape `[batch_size, k]`.
      n: an integer.
    Returns:
      A `Tensor` with the same type as `values` and shape `[batch_size, n]`.
    """
    batch, k = tf.unstack(tf.shape(indices), num=2)
    indices_flat = tf.reshape(indices, [-1]) + tf.div(tf.range(batch * k), k) * n
    ret_flat = tf.unsorted_segment_sum(
        tf.reshape(values, [-1]), indices_flat, batch * n)
    return tf.reshape(ret_flat, [batch, n])