code: def find_bins(
      x,
      edges,
      extend_lower_interval=False,
      extend_upper_interval=False,
      dtype=None,
      name=None,
  ):
      """Bin values into discrete intervals.
  
      Given `edges = [c0, ..., cK]`, defining intervals
      `I_0 = [c0, c1)`, `I_1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
      This function returns `bins`, such that `x[i]` lies within `I_{bins[i]}`.
  
      Args:
        x:  Numeric `N-D` `Tensor` with `N > 0`.
        edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
          of intervals.  Must either be `1-D` or have
          `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
          designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
          corresponding dimensions of `x`.
        extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
          interval `I0` to `(-inf, c1]`.
        extend_upper_interval:  Python `bool`.  If `True`, extend the upper
          interval `I_{K-1}` to `[c_{K-1}, +inf)`.
        dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
          This effects the output values when `x` is below/above the intervals,
          which will be `-1/K+1` for `int` types and `NaN` for `float`s.
          At indices where `x` is `NaN`, the output values will be `0` for `int`
          types and `NaN` for floats.
        name:  A Python string name to prepend to created ops. Default: 'find_bins'
  
      Returns:
        bins: `Tensor` with same `shape` as `x` and `dtype`.
          Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
          `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.
  
      Raises:
        ValueError:  If `edges.shape[0]` is determined to be less than 2.
  
      #### Examples
  
      Cut a `1-D` array
  
      ```python
      x = [0., 5., 6., 10., 20.]
      edges = [0., 5., 10.]
      tfp.stats.find_bins(x, edges)
      ==> [0., 1., 1., 1., np.nan]
      ```
  
      Cut `x` into its deciles
  
      ```python
      x = tf.random.uniform(shape=(100, 200))
      decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
      bins = tfp.stats.find_bins(x, edges=decile_edges)
      bins.shape
      ==> (100, 200)
      tf.reduce_mean(bins == 0.)
      ==> approximately 0.1
      tf.reduce_mean(bins == 1.)
      ==> approximately 0.1
      ```
  
      """
      # TFP users may be surprised to see the "action" in the leftmost dim of
      # edges, rather than the rightmost (event) dim.  Why?
      # 1. Most likely you created edges by getting quantiles over samples, and
      #    quantile/percentile return these edges in the leftmost (sample) dim.
      # 2. Say you have event_shape = [5], then we expect the bin will be different
      #    for all 5 events, so the index of the bin should not be in the event dim.
      with tf.name_scope(name or "find_bins"):
          in_type = dtype_util.common_dtype(
              [x, edges], dtype_hint=tf.float32
          )
          edges = tf.convert_to_tensor(
              edges, name="edges", dtype=in_type
          )
          x = tf.convert_to_tensor(x, name="x", dtype=in_type)
  
          if (
              tf.compat.dimension_value(edges.shape[0]) is not None
              and tf.compat.dimension_value(edges.shape[0]) < 2
          ):
              raise ValueError(
                  "First dimension of `edges` must have length > 1 to index 1 or "
                  "more bin. Found: {}".format(edges.shape)
              )
  
          flattening_x = (
              tensorshape_util.rank(edges.shape) == 1
              and tensorshape_util.rank(x.shape) > 1
          )
  
          if flattening_x:
              x_orig_shape = ps.shape(x)
              x = tf.reshape(x, [-1])
  
          if dtype is None:
              dtype = in_type
          dtype = tf.as_dtype(dtype)
  
          # Move first dims into the rightmost.
          x_permed = distribution_util.rotate_transpose(x, shift=-1)
          edges_permed = distribution_util.rotate_transpose(
              edges, shift=-1
          )
  
          # If...
          #   x_permed = [0, 1, 6., 10]
          #   edges = [0, 5, 10.]
          #   ==> almost_output = [0, 1, 2, 2]
          searchsorted_type = (
              dtype if dtype in [tf.int32, tf.int64] else None
          )
          almost_output_permed = tf.searchsorted(
              sorted_sequence=edges_permed,
              values=x_permed,
              side="right",
              out_type=searchsorted_type,
          )
          # Move the rightmost dims back to the leftmost.
          almost_output = tf.cast(
              distribution_util.rotate_transpose(
                  almost_output_permed, shift=1
              ),
              dtype,
          )
  
          # In above example, we want [0, 0, 1, 1], so correct this here.
          bins = tf.clip_by_value(
              almost_output - 1,
              tf.cast(0, dtype),
              tf.cast(tf.shape(edges)[0] - 2, dtype),
          )
  
          if not extend_lower_interval:
              low_fill = np.nan if dtype_util.is_floating(dtype) else -1
              bins = tf.where(
                  x < tf.expand_dims(edges[0], 0),
                  tf.cast(low_fill, dtype),
                  bins,
              )
  
          if not extend_upper_interval:
              up_fill = (
                  np.nan
                  if dtype_util.is_floating(dtype)
                  else tf.shape(edges)[0] - 1
              )
              bins = tf.where(
                  x > tf.expand_dims(edges[-1], 0),
                  tf.cast(up_fill, dtype),
                  bins,
              )
  
          if flattening_x:
              bins = tf.reshape(bins, x_orig_shape)
  
          return bins
  
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 2 for calling a global function with example usage: find_bins(x, edges, extend_lower_interval, extend_upper_interval, dtype, name) and returns: bins