code: def _log_loosum_exp_impl(logx, axis, keepdims, compute_mean):
      """Implementation for `*loosum*` functions."""
      with tf.name_scope("log_loosum_exp_impl"):
          logx = tf.convert_to_tensor(logx, name="logx")
          dtype = dtype_util.as_numpy_dtype(logx.dtype)
  
          if axis is not None:
              x = np.array(axis)
              axis = (
                  tf.convert_to_tensor(
                      axis, name="axis", dtype_hint=tf.int32
                  )
                  if x.dtype is np.object_
                  else x.astype(np.int32)
              )
  
          log_sum_x = tf.reduce_logsumexp(
              logx, axis=axis, keepdims=True
          )
  
          # Later we'll want to compute the mean from a sum so we calculate the number
          # of reduced elements, n.
          n = ps.size(logx) // ps.size(log_sum_x)
          n = ps.cast(n, dtype)
  
          # log_loosum_x[i] =
          # = logsumexp(logx[j] : j != i)
          # = log( exp(logsumexp(logx)) - exp(logx[i]) )
          # = log( exp(logsumexp(logx - logx[i])) exp(logx[i])  - exp(logx[i]))
          # = logx[i] + log(exp(logsumexp(logx - logx[i])) - 1)
          # = logx[i] + log(exp(logsumexp(logx) - logx[i]) - 1)
          # = logx[i] + softplus_inverse(logsumexp(logx) - logx[i])
          d = log_sum_x - logx
          # We use `d != 0` rather than `d > 0.` because `d < 0.` should never happen;
          # if it does we want to complain loudly (which `softplus_inverse` will).
          d_ok = tf.not_equal(d, 0.0)
          safe_d = tf.where(d_ok, d, 1.0)
          d_ok_result = logx + softplus_inverse(safe_d)
  
          neg_inf = tf.constant(-np.inf, dtype=dtype)
  
          # When not(d_ok) and is_largest then we manually compute the
          # log_loosum_x. (We can efficiently do this for any one point but not all,
          # hence we still need the above calculation.) This is good because when
          # this condition is met, we cannot use the above calculation; its -inf.
          # We now compute the log-leave-out-max-sum, replicate it to every
          # point and make sure to select it only when we need to.
          max_logx = tf.reduce_max(logx, axis=axis, keepdims=True)
          is_largest = tf.equal(logx, max_logx)
          log_lomsum_x = tf.reduce_logsumexp(
              tf.where(is_largest, neg_inf, logx),
              axis=axis,
              keepdims=True,
          )
          d_not_ok_result = tf.where(is_largest, log_lomsum_x, neg_inf)
  
          log_loosum_x = tf.where(d_ok, d_ok_result, d_not_ok_result)
  
          # We now squeeze log_sum_x so as if we used `keepdims=False`.
          # TODO(b/136176077): These mental gymnastics could all be replaced with
          # `tf.squeeze(log_sum_x, axis)` if tf.squeeze supported Tensor valued `axis`
          # arguments.
          if not keepdims:
              if axis is None:
                  keepdims = np.array([], dtype=np.int32)
              else:
                  rank = ps.rank(logx)
                  keepdims = ps.setdiff1d(
                      ps.range(rank), ps.non_negative_axis(axis, rank)
                  )
              squeeze_shape = tf.gather(
                  ps.shape(logx), indices=keepdims
              )
              log_sum_x = tf.reshape(log_sum_x, shape=squeeze_shape)
              if ps.is_numpy(keepdims):
                  tensorshape_util.set_shape(
                      log_sum_x, np.array(logx.shape)[keepdims]
                  )
  
          # Set static shapes just in case we lost them.
          tensorshape_util.set_shape(n, [])
          tensorshape_util.set_shape(log_loosum_x, logx.shape)
  
          if not compute_mean:
              return log_loosum_x, log_sum_x, n
  
          log_nm1 = ps.log(max(1.0, n - 1.0))
          log_n = ps.log(n)
          return log_loosum_x - log_nm1, log_sum_x - log_n, n
  
file_path: probability/tensorflow_probability/python/stats/leave_one_out.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.generic import log_add_exp
  from tensorflow_probability.python.math.generic import softplus_inverse
label_desc: Use image node_id 6 for calling a global function with example usage: _log_loosum_exp_impl(logx, axis, keepdims, compute_mean) and returns: unknown,  unknown,  n, log_loosum_x,  log_sum_x,  n