code: def _find_bins(x, edges, axis, dtype=tf.int64, name=None):
      """Like `tfp.stats.find_bins` but correctly handles quantiles axis arg."""
      with tf.name_scope(name or "find_bins"):
          # We can't do this:
          #   return tf.cast(quantiles_lib.find_bins(x, edges=edges), dtype=tf.int64)
          # because it doesn't seem to correctly handle axis!=-1. This is a bug in TFP
          # and should be fixed. Furthermore, the following is probably more efficient
          # than tfp.stats..find_bins anyway.
          num_buckets = ps.size0(edges) - 1
          # First, we need to have `keepdims=True` semantics for edges.
          axis = axis % ps.rank(x)
          edges = tf.expand_dims(edges, axis + 1)
          # We now find the bucket which is is the "first larger", then subtract one
          # to get the bucket which is the "last smaller". Care must be taken for the
          # max element.
          pred = x < edges
          # The following is equivalent to:
          #    tf.argmin(tf.cast(~pred, dtype), axis=0))
          # yet gives the same implementation across TPU/GPU/CPU.
          # As a bonus, we can also leverage the `sorted=True` behavior.
          _, bucket_larger = tf.math.top_k(
              tf.cast(
                  tf.transpose(
                      pred,
                      ps.pad(
                          ps.range(1, ps.rank(pred)), paddings=[[0, 1]]
                      ),
                  ),
                  dtype,
              ),
              k=1,
              sorted=True,
          )
          bucket_larger = bucket_larger[..., 0]
  
          bucket_larger = tf.where(
              pred[-1],  # == ~tf.math.reduce_all(pred, axis=0)
              tf.cast(bucket_larger, dtype),
              tf.cast(num_buckets, dtype),
          )
          return bucket_larger - 1
  
file_path: probability/tensorflow_probability/python/stats/calibration.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.stats import quantiles
label_desc: Use image node_id 8 for calling a global function with example usage: _find_bins(x, edges, axis, dtype, name) and returns: unknown