code: def estimate_tails(func, target, shape, dtype):
      """Estimates approximate tail quantiles.
  
      This runs a simple Adam iteration to determine tail quantiles. The
      objective is to find an `x` such that:
      ```
      func(x) == target
      ```
      For instance, if `func` is a CDF and the target is a quantile value, this
      would find the approximate location of that quantile. Note that `func` is
      assumed to be monotonic. When each tail estimate has passed the optimal value
      of `x`, the algorithm does 100 additional iterations and then stops.
  
      This operation is vectorized. The tensor shape of `x` is given by `shape`, and
      `target` must have a shape that is broadcastable to the output of `func(x)`.
  
      Args:
        func: A callable that computes cumulative distribution function, survival
          function, or similar.
        target: The desired target value.
        shape: The shape of the `tf.Tensor` representing `x`.
        dtype: The `tf.dtypes.Dtype` of the computation (and the return value).
  
      Returns:
        A `tf.Tensor` representing the solution (`x`).
      """
      with tf.name_scope("estimate_tails"):
          dtype = tf.as_dtype(dtype)
          shape = tf.convert_to_tensor(shape, tf.int32)
          target = tf.convert_to_tensor(target, dtype)
  
          def loop_cond(
              tails, m, v, loss, count, best_tails, best_loss
          ):
              del tails, m, v, best_tails, best_loss  # unused
              # By checking `loss`, we catch NaNs and protect against infinite loops
              # from perfect initial guesses where there is no zero-crossing.
              return tf.logical_and(
                  tf.reduce_max(loss) > 1e-8, tf.reduce_min(count) < 100
              )
  
          def loop_body(
              tails, prev_m, prev_v, loss, count, best_tails, best_loss
          ):
              del loss  # always recomputed
              with tf.GradientTape(
                  watch_accessed_variables=False
              ) as tape:
                  tape.watch(tails)
                  loss = abs(func(tails) - target)
  
              # Keep track of the best (lowest loss) value so far.
              condition = loss < best_loss
              best_tails = tf.where(condition, tails, best_tails)
              best_loss = tf.where(condition, loss, best_loss)
  
              grad = tape.gradient(loss, tails)
              m = (prev_m + grad) / 2  # Adam mean estimate.
              v = (
                  prev_v + tf.square(grad)
              ) / 2  # Adam variance estimate.
  
              # Reduce learning rate as count increases. This should lead to a more
              # accurate final value.
              k = tf.math.sqrt(tf.cast(count + 1, m.dtype))
              tails -= 0.1 * m / (k * tf.sqrt(v) + 1e-20)
  
              # Start counting when the gradient flips sign. Since the function is
              # monotonic, m must have the same sign in all initial iterations, until
              # the optimal point is crossed. At that point the gradient flips sign.
              count = tf.where(
                  tf.math.logical_or(count > 0, prev_m * grad < 0),
                  count + 1,
                  count,
              )
              return tails, m, v, loss, count, best_tails, best_loss
  
          init_tails = tf.zeros(shape, dtype=dtype)
          init_m = tf.zeros(shape, dtype=dtype)
          init_v = tf.ones(shape, dtype=dtype)
          init_loss = init_v * dtype.max
          init_count = tf.zeros(shape, dtype=tf.int32)
          loop_vars = (
              init_tails,
              init_m,
              init_v,
              init_loss,
              init_count,
              init_tails,
              init_loss,
          )
          return tf.while_loop(loop_cond, loop_body, loop_vars)[-2]
  
imports: import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: estimate_tails(func, target, shape, dtype) and returns: tf, tf, tails,  m,  v,  loss,  count,  best_tails,  best_loss