code: def _log_softmax(x, axis, name=None):
      """Alternative to `tf.log_softmax` which correctly supports axis arg."""
      with tf.name_scope(name or "log_softmax"):
          if axis is None:
              return tf.math.log_softmax(x, axis=None, name=name)
          rank = ps.rank(axis)
          if rank == 0:
              return tf.math.log_softmax(x, axis=axis, name=name)
          if rank == 1:
              return tf.math.log_softmax(x, axis=axis[0], name=name)
          # The following handles the case when axis is a vector and which is not
          # currently supported by tf.math.log_softmax.
          x = tf.convert_to_tensor(x, dtype_hint=tf.float32, name="x")
          return x - tf.reduce_logsumexp(x, axis=axis, keepdims=True)
  
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
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
label_desc: Use image node_id 13 for calling a global function with example usage: _log_softmax(x, axis, name) and returns: unknown, tf, tf, tf