code: 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.name_scope(name or "count_integers"):
          if axis is None:
              return tf.math.bincount(
                  arr,
                  weights=weights,
                  minlength=minlength,
                  maxlength=maxlength,
                  dtype=dtype,
              )
  
          arr = tf.convert_to_tensor(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
          )
          minlength = (
              minlength
              if minlength is not None
              else tf.reduce_max(arr) + 1
          )
          maxlength = (
              maxlength
              if maxlength is not None
              else tf.reduce_max(arr) + 1
          )
  
          # 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,
                  fn_output_signature=dtype,
              )
          else:
              weights = tf.convert_to_tensor(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],
                  fn_output_signature=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 = ps.gather(ps.shape(arr), indices=not_axis)
  
          # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
          out_shape = ps.concat([[-1], not_axis_shape], axis=0)
  
          return tf.reshape(flat_counts_t, out_shape)
  
file_path: probability/tensorflow_probability/python/stats/quantiles.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: count_integers(arr, weights, minlength, maxlength, axis, dtype, name) and returns: tf, tf, tf, tf, tf