code: def _softmax(logits):
      """Applies softmax non-linearity on inputs."""
      return np.exp(logits) / np.sum(
          np.exp(logits), axis=-1, keepdims=True
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  from trfl import vtrace_ops
label_desc: Use image node_id 2 for calling a global function with example usage: _softmax(logits) and returns: unknown