def _new_unsorted_segment_mean(data, segment_ids, num_segments, name=None): """ERROR: docstring should have been added programatically. """ with ops.name_scope( name, "UnsortedSegmentMean", [data, segment_ids, num_segments]) as name: # Note that data can be a vector-like list (or an n-dimensional # tensor-like list of lists). We convert to tensor here to replicate the # behavior of the pre-existing op. data = tf.convert_to_tensor(data) N = math_ops._unsorted_segment_N(data, segment_ids, num_segments) # Note that this patch does not provide determinism when the dtype of the # data argument is tf.float64 or tf.complex128. orig_dtype = data.dtype if 'float' in str(orig_dtype): data = tf.cast(data, dtype=tf.float64) elif 'complex' in str(orig_dtype): data = tf.cast(data, dtype=tf.complex128) if not context.executing_eagerly(): data = ops.convert_to_tensor(data, name="input_data") segment_ids = ops.convert_to_tensor(segment_ids, name="segment_ids") num_segments = ops.convert_to_tensor(num_segments, name="num_segments") result = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments) return tf.cast(result, dtype=orig_dtype) / N
def unsorted_segment_mean(data, segment_ids, num_segments, name=None): with ops.name_scope(name, "UnsortedSegmentMean"): data = ops.convert_to_tensor(data) segment_ids = ops.convert_to_tensor(segment_ids) N = _unsorted_segment_N(data, segment_ids, num_segments) summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments) return summed / N
def _unsorted_segment_N(data, segment_ids, num_segments): """ Helper function for unsorted_segment_mean/_sqrtN. Computes the number of segment entries with 0-entries set to 1 to allow division by N. """ # bincount doesn't support negative indices so we use unsorted_segment_sum segment_ids_shape = array_ops.shape_internal(segment_ids) ones_tensor = array_ops.ones(segment_ids_shape, dtype=data.dtype) N = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments) # add dimensions for all non-reduced axes ndims_output = data.shape.ndims - segment_ids.shape.ndims broadcast_shape = [num_segments] + [1] * ndims_output N = array_ops.reshape(N, broadcast_shape) return gen_math_ops.maximum(N, 1)
def _UnsortedSegmentProdGrad(op, grad): """ Gradient for UnsortedSegmentProd. The gradient can be expressed for each segment by dividing the segment's product by each element of the segment input tensor, but this approach can't deal with zeros in the input. Unlike reduce_prod we can't use cumsum here as individual segments may have a different number of elements. Therefore we consider three cases: 1) A segment input contains no zeros and we can safely divide by the input tensor. 2) A segment contains exactly one zero. Then the gradient of each input of the segment is zero except for the 0-input, there the gradient is the product of the remaining segment entries. 3) A segment contains at least two zeros. The gradient is zero for all segment inputs. """ # Note that unsorted_segment_sum will filter out the negative indices, # so we don't need to do a logical_and with is_positive here is_zero = math_ops.equal(op.inputs[0], 0) num_zeros = gen_math_ops.unsorted_segment_sum( math_ops.cast(is_zero, dtype=dtypes.int32), op.inputs[1], op.inputs[2]) # handle case 3 and set the gradient to 0 for segments with more than one # 0 as input grad = array_ops.where(math_ops.greater(num_zeros, 1), array_ops.zeros_like(grad), grad) # replace all zeros with ones and compute the unsorted_segment_prod non_zero_data = array_ops.where(is_zero, array_ops.ones_like(op.inputs[0]), op.inputs[0]) non_zero_prod = gen_math_ops.unsorted_segment_prod(non_zero_data, op.inputs[1], op.inputs[2]) # clip the indices for gather to be positive zero_clipped_indices = math_ops.maximum(op.inputs[1], array_ops.zeros_like(op.inputs[1])) gathered_prod = array_ops.gather(op.outputs[0], zero_clipped_indices) gathered_non_zero_prod = array_ops.gather(non_zero_prod, zero_clipped_indices) prod_divided_by_el = gathered_prod / op.inputs[0] # May contain nan/inf. # Now fetch the individual results for segments containing 0 and those that # don't. is_zero will also fetch results for entries with negative index # but the following gather_drop_negatives sets the corresponding entry in # grad to 0 for these partial_derivative = array_ops.where(is_zero, gathered_non_zero_prod, prod_divided_by_el) gathered_grad = _GatherDropNegatives(grad, op.inputs[1], zero_clipped_indices)[0] return gathered_grad * partial_derivative, None, None
def _UnsortedSegmentProdGrad(op, grad): """ Gradient for UnsortedSegmentProd. The gradient can be expressed for each segment by dividing the segment's product by each element of the segment input tensor, but this approach can't deal with zeros in the input. Unlike reduce_prod we can't use cumsum here as individual segments may have a different number of elements. Therefore we consider three cases: 1) A segment input contains no zeros and we can safely divide by the input tensor. 2) A segment contains exactly one zero. Then the gradient of each input of the segment is zero except for the 0-input, there the gradient is the product of the remaining segment entries. 3) A segment contains at least two zeros. The gradient is zero for all segment inputs. """ # Note that unsorted_segment_sum will filter out the negative indices, # so we don't need to do a logical_and with is_positive here is_zero = math_ops.equal(op.inputs[0], 0) num_zeros = gen_math_ops.unsorted_segment_sum( math_ops.cast(is_zero, dtype=dtypes.int32), op.inputs[1], op.inputs[2]) # handle case 3 and set the gradient to 0 for segments with more than one # 0 as input grad = array_ops.where(math_ops.greater(num_zeros, 1), array_ops.zeros_like(grad), grad) # replace all zeros with ones and compute the unsorted_segment_prod non_zero_data = array_ops.where(is_zero, array_ops.ones_like(op.inputs[0]), op.inputs[0]) non_zero_prod = gen_math_ops.unsorted_segment_prod( non_zero_data, op.inputs[1], op.inputs[2]) # clip the indices for gather to be positive zero_clipped_indices = math_ops.maximum(op.inputs[1], array_ops.zeros_like(op.inputs[1])) gathered_prod = array_ops.gather(op.outputs[0], zero_clipped_indices) gathered_non_zero_prod = array_ops.gather(non_zero_prod, zero_clipped_indices) prod_divided_by_el = gathered_prod / op.inputs[0] # May contain nan/inf. # Now fetch the individual results for segments containing 0 and those that # don't. is_zero will also fetch results for entries with negative index # but the following gather_drop_negatives sets the corresponding entry in # grad to 0 for these partial_derivative = array_ops.where(is_zero, gathered_non_zero_prod, prod_divided_by_el) gathered_grad = _GatherDropNegatives(grad, op.inputs[1], zero_clipped_indices)[0] return gathered_grad * partial_derivative, None, None
def bincount(arr, weights=None, minlength=None, maxlength=None, dtype=dtypes.int32, name=None, axis=None, binary_output=False): """Counts the number of occurrences of each value in an integer array. If `minlength` and `maxlength` are not given, returns a vector with length `tf.reduce_max(arr) + 1` if `arr` is non-empty, and length 0 otherwise. If `weights` are non-None, then index `i` of the output stores the sum of the value in `weights` at each index where the corresponding value in `arr` is `i`. ```python values = tf.constant([1,1,2,3,2,4,4,5]) tf.math.bincount(values) #[0 2 2 1 2 1] ``` Vector length = Maximum element in vector `values` is 5. Adding 1, which is 6 will be the vector length. Each bin value in the output indicates number of occurrences of the particular index. Here, index 1 in output has a value 2. This indicates value 1 occurs two times in `values`. ```python values = tf.constant([1,1,2,3,2,4,4,5]) weights = tf.constant([1,5,0,1,0,5,4,5]) tf.math.bincount(values, weights=weights) #[0 6 0 1 9 5] ``` Bin will be incremented by the corresponding weight instead of 1. Here, index 1 in output has a value 6. This is the summation of weights corresponding to the value in `values`. **Bin-counting on a certain axis** This example takes a 2 dimensional input and returns a `Tensor` with bincounting on each sample. >>> data = np.array([[1, 2, 3, 0], [0, 0, 1, 2]], dtype=np.int32) >>> tf.math.bincount(data, axis=-1) <tf.Tensor: shape=(2, 4), dtype=int32, numpy= array([[1, 1, 1, 1], [2, 1, 1, 0]], dtype=int32)> **Bin-counting with binary_output** This example gives binary output instead of counting the occurrence. >>> data = np.array([[1, 2, 3, 0], [0, 0, 1, 2]], dtype=np.int32) >>> tf.math.bincount(data, axis=-1, binary_output=True) <tf.Tensor: shape=(2, 4), dtype=int32, numpy= array([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=int32)> Args: arr: A Tensor, RaggedTensor, or SparseTensor whose values should be counted. These tensors must have a rank of 2 if `axis=-1`. weights: If non-None, must be the same shape as arr. For each value in `arr`, the bin will be incremented by the corresponding weight instead of 1. minlength: If given, ensures the output has length at least `minlength`, padding with zeros at the end if necessary. maxlength: If given, skips values in `arr` that are equal or greater than `maxlength`, ensuring that the output has length at most `maxlength`. dtype: If `weights` is None, determines the type of the output bins. name: A name scope for the associated operations (optional). axis: The axis to slice over. Axes at and below `axis` will be flattened before bin counting. Currently, only `0`, and `-1` are supported. If None, all axes will be flattened (identical to passing `0`). binary_output: If True, this op will output 1 instead of the number of times a token appears (equivalent to one_hot + reduce_any instead of one_hot + reduce_add). Defaults to False. Returns: A vector with the same dtype as `weights` or the given `dtype`. The bin values. Raises: `InvalidArgumentError` if negative values are provided as an input. """ name = "bincount" if name is None else name with ops.name_scope(name): # Somehow forward compatible needs to be False. if not binary_output and axis is None: arr = ops.convert_to_tensor(arr, name="arr", dtype=dtypes.int32) array_is_nonempty = math_ops.reduce_prod(array_ops.shape(arr)) > 0 output_size = math_ops.cast(array_is_nonempty, dtypes.int32) * ( math_ops.reduce_max(arr) + 1) if minlength is not None: minlength = ops.convert_to_tensor(minlength, name="minlength", dtype=dtypes.int32) output_size = gen_math_ops.maximum(minlength, output_size) if maxlength is not None: maxlength = ops.convert_to_tensor(maxlength, name="maxlength", dtype=dtypes.int32) output_size = gen_math_ops.minimum(maxlength, output_size) if weights is not None: weights = ops.convert_to_tensor(weights, name="weights") return gen_math_ops.unsorted_segment_sum( weights, arr, output_size) weights = constant_op.constant([], dtype) arr = array_ops.reshape(arr, [-1]) return gen_math_ops.bincount(arr, output_size, weights) if not isinstance(arr, sparse_tensor.SparseTensor): arr = ragged_tensor.convert_to_tensor_or_ragged_tensor(arr, name="arr") if weights is not None: if not isinstance(weights, sparse_tensor.SparseTensor): weights = ragged_tensor.convert_to_tensor_or_ragged_tensor( weights, name="weights") if weights is not None and binary_output: raise ValueError( "Arguments `binary_output` and `weights` are mutually " "exclusive. Please specify only one.") if not arr.dtype.is_integer: arr = math_ops.cast(arr, dtypes.int32) if axis is None: axis = 0 if axis not in [0, -1]: raise ValueError( f"Unsupported value for argument axis={axis}. Only 0 and" " -1 are currently supported.") if isinstance(arr, ragged_tensor.RaggedTensor): array_is_nonempty = math_ops.reduce_prod( array_ops.shape(arr.values)) > 0 else: array_is_nonempty = math_ops.reduce_prod(array_ops.shape(arr)) > 0 if isinstance(arr, sparse_tensor.SparseTensor): output_size = math_ops.cast(array_is_nonempty, arr.dtype) * ( math_ops.reduce_max(arr.values) + 1) else: output_size = math_ops.cast( array_is_nonempty, arr.dtype) * (math_ops.reduce_max(arr) + 1) if minlength is not None: minlength = ops.convert_to_tensor(minlength, name="minlength", dtype=arr.dtype) output_size = gen_math_ops.maximum(minlength, output_size) if maxlength is not None: maxlength = ops.convert_to_tensor(maxlength, name="maxlength", dtype=arr.dtype) output_size = gen_math_ops.minimum(maxlength, output_size) if axis == 0: if isinstance(arr, sparse_tensor.SparseTensor): if weights is not None: weights = validate_sparse_weights(arr, weights, dtype) arr = arr.values elif isinstance(arr, ragged_tensor.RaggedTensor): if weights is not None: weights = validate_ragged_weights(arr, weights, dtype) arr = arr.values else: if weights is not None: weights = array_ops.reshape(weights, [-1]) arr = array_ops.reshape(arr, [-1]) if isinstance(arr, sparse_tensor.SparseTensor): weights = validate_sparse_weights(arr, weights, dtype) return gen_math_ops.sparse_bincount(indices=arr.indices, values=arr.values, dense_shape=arr.dense_shape, size=output_size, weights=weights, binary_output=binary_output) elif isinstance(arr, ragged_tensor.RaggedTensor): weights = validate_ragged_weights(arr, weights, dtype) return gen_math_ops.ragged_bincount(splits=arr.row_splits, values=arr.values, size=output_size, weights=weights, binary_output=binary_output) else: weights = validate_dense_weights(arr, weights, dtype) return gen_math_ops.dense_bincount(input=arr, size=output_size, weights=weights, binary_output=binary_output)