tensorflow_probability/python/stats/quantiles.py
def count_integers(arr,
                   weights=None,
                   minlength=None,
                   maxlength=None,
                   axis=None,
                   dtype=tf.int32,
                   name=None):
  """Counts the number of occurrences of each value in an integer array `arr`.

  Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
  dimensions to reduce over.  With
    `~axis = [i for i in range(arr.ndim) if i not in axis]`,
  this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.

  If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
  if `arr` is non-empty, and 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`.

  Args:
    arr: An `int32` `Tensor` of non-negative values.
    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`.
    axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
      dimensions in `arr` to reduce over.
      `Default value:` `None`, meaning reduce over all dimensions.
    dtype: If `weights` is None, determines the type of the output bins.
    name: A name scope for the associated operations (optional).

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
  """
  with tf.compat.v1.name_scope(
      name, 'count_integers', values=[arr, weights, minlength, maxlength,
                                      axis]):
    if axis is None:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='arr')
    arr_ndims = _get_static_ndims(arr, expect_static=True)

    axis = _make_static_axis_non_negative_list(axis, arr_ndims)

    # ~axis from docstring.  Dims in arr that are not in axis.
    not_axis = sorted(set(range(arr_ndims)).difference(axis))

    # If we're reducing over everything, just use standard bincount.
    if not not_axis:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    # Move dims in ~axis to the left, so we can tf.map_fn bincount over them,
    # Producing counts for every index I in ~axis.
    # Thus, flat_arr is not totally flat, it just has the dims in ~axis
    # flattened.
    flat_arr = _move_dims_to_flat_end(arr, not_axis, arr_ndims, right_end=False)

    # tf.map_fn over dim 0.
    if weights is None:

      def one_bincount(arr_slice):
        return tf.math.bincount(
            arr_slice,
            weights=None,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)
    else:
      weights = tf.convert_to_tensor(value=weights, name='weights')
      _get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)
      flat_weights = _move_dims_to_flat_end(
          weights, not_axis, arr_ndims, right_end=False)

      def one_bincount(arr_and_weights_slices):
        arr_slice, weights_slice = arr_and_weights_slices
        return tf.math.bincount(
            arr_slice,
            weights=weights_slice,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(
          one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)

    # flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0.
    # bincount needs to have the K bins in axis 0, so transpose...
    flat_counts_t = tf.transpose(a=flat_counts, perm=[1, 0])

    # Throw in this assert, to ensure shape assumptions are correct.
    _get_static_ndims(flat_counts_t, expect_ndims=2, expect_static=True)

    # not_axis_shape = arr.shape[~axis]
    not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    out_shape = tf.concat([[-1], not_axis_shape], axis=0)

    return tf.reshape(flat_counts_t, out_shape)
def count_integers(arr,
                   weights=None,
                   minlength=None,
                   maxlength=None,
                   axis=None,
                   dtype=tf.int32,
                   name=None):
  with tf.compat.v1.name_scope(
      name, 'count_integers', values=[arr, weights, minlength, maxlength,
                                      axis]):
    if axis is None:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='arr')
    arr_ndims = _get_static_ndims(arr, expect_static=True)

    axis = _make_static_axis_non_negative_list(axis, arr_ndims)

    # ~axis from docstring.  Dims in arr that are not in axis.
    not_axis = sorted(set(range(arr_ndims)).difference(axis))

    # If we're reducing over everything, just use standard bincount.
    if not not_axis:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    # Move dims in ~axis to the left, so we can tf.map_fn bincount over them,
    # Producing counts for every index I in ~axis.
    # Thus, flat_arr is not totally flat, it just has the dims in ~axis
    # flattened.
    flat_arr = _move_dims_to_flat_end(arr, not_axis, arr_ndims, right_end=False)

    # tf.map_fn over dim 0.
    if weights is None:

      def one_bincount(arr_slice):
        return tf.math.bincount(
            arr_slice,
            weights=None,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)
    else:
      weights = tf.convert_to_tensor(value=weights, name='weights')
      _get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)
      flat_weights = _move_dims_to_flat_end(
          weights, not_axis, arr_ndims, right_end=False)

      def one_bincount(arr_and_weights_slices):
        arr_slice, weights_slice = arr_and_weights_slices
        return tf.math.bincount(
            arr_slice,
            weights=weights_slice,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(
          one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)

    # flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0.
    # bincount needs to have the K bins in axis 0, so transpose...
    flat_counts_t = tf.transpose(a=flat_counts, perm=[1, 0])

    # Throw in this assert, to ensure shape assumptions are correct.
    _get_static_ndims(flat_counts_t, expect_ndims=2, expect_static=True)

    # not_axis_shape = arr.shape[~axis]
    not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    out_shape = tf.concat([[-1], not_axis_shape], axis=0)

    return tf.reshape(flat_counts_t, out_shape)
def count_integers ( arr , weights = None , minlength = None , maxlength = None , axis = None , dtype = tf . int32 , name = None ) : with tf . compat . v1 . name_scope ( name , 'count_integers' , values = [ arr , weights , minlength , maxlength , axis ] ) : if axis is None : return tf . math . bincount ( arr , weights = weights , minlength = minlength , maxlength = maxlength , dtype = dtype ) arr = tf . convert_to_tensor ( value = arr , dtype = tf . int32 , name = 'arr' ) arr_ndims = _get_static_ndims ( arr , expect_static = True ) axis = _make_static_axis_non_negative_list ( axis , arr_ndims ) # ~axis from docstring.  Dims in arr that are not in axis. not_axis = sorted ( set ( range ( arr_ndims ) ) . difference ( axis ) ) # If we're reducing over everything, just use standard bincount. if not not_axis : return tf . math . bincount ( arr , weights = weights , minlength = minlength , maxlength = maxlength , dtype = dtype ) # Move dims in ~axis to the left, so we can tf.map_fn bincount over them, # Producing counts for every index I in ~axis. # Thus, flat_arr is not totally flat, it just has the dims in ~axis # flattened. flat_arr = _move_dims_to_flat_end ( arr , not_axis , arr_ndims , right_end = False ) # tf.map_fn over dim 0. if weights is None : def one_bincount ( arr_slice ) : return tf . math . bincount ( arr_slice , weights = None , minlength = minlength , maxlength = maxlength , dtype = dtype ) flat_counts = tf . map_fn ( one_bincount , elems = flat_arr , dtype = dtype ) else : weights = tf . convert_to_tensor ( value = weights , name = 'weights' ) _get_static_ndims ( weights , expect_static = True , expect_ndims = arr_ndims ) flat_weights = _move_dims_to_flat_end ( weights , not_axis , arr_ndims , right_end = False ) def one_bincount ( arr_and_weights_slices ) : arr_slice , weights_slice = arr_and_weights_slices return tf . math . bincount ( arr_slice , weights = weights_slice , minlength = minlength , maxlength = maxlength , dtype = dtype ) flat_counts = tf . map_fn ( one_bincount , elems = [ flat_arr , flat_weights ] , dtype = weights . dtype ) # flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0. # bincount needs to have the K bins in axis 0, so transpose... flat_counts_t = tf . transpose ( a = flat_counts , perm = [ 1 , 0 ] ) # Throw in this assert, to ensure shape assumptions are correct. _get_static_ndims ( flat_counts_t , expect_ndims = 2 , expect_static = True ) # not_axis_shape = arr.shape[~axis] not_axis_shape = tf . gather ( tf . shape ( input = arr ) , indices = not_axis ) # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis out_shape = tf . concat ( [ [ - 1 ] , not_axis_shape ] , axis = 0 ) return tf . reshape ( flat_counts_t , out_shape )
Counts the number of occurrences of each value in an integer array `arr`.

  Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
  dimensions to reduce over.  With
    `~axis = [i for i in range(arr.ndim) if i not in axis]`,
  this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.

  If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
  if `arr` is non-empty, and 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`.

  Args:
    arr: An `int32` `Tensor` of non-negative values.
    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`.
    axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
      dimensions in `arr` to reduce over.
      `Default value:` `None`, meaning reduce over all dimensions.
    dtype: If `weights` is None, determines the type of the output bins.
    name: A name scope for the associated operations (optional).

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
Counts the number of occurrences of each value in an integer array arr .