code: def log_softmax(x, axis=-1):  # pylint: disable=invalid-name
      """Transforms activation vectors to log-probability vectors.
  
      Log probability vectors are derived by, in effect, applying softmax to raw
      activation vectors and then applying log element-wise. The actual
      implementation uses a mathematically valid simplification of this.
  
      Args:
        x: An ndarray with activation vectors along the given axis.
        axis: Axis along which values are grouped for computing log softmax.
  
      Returns:
        An ndarray containing log-probability vectors derived from the raw
        activation vectors in `x`.
      """
      return x - fastmath.logsumexp(x, axis=axis, keepdims=True)
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 16 for calling a global function with example usage: log_softmax(x, axis) and returns: unknown