code: def softmax(explainer, op, *grads):
      """Just decompose softmax into its components and recurse, we can handle all of them :)
  
      We assume the 'axis' is the last dimension because the TF codebase swaps the 'axis' to
      the last dimension before the softmax op if 'axis' is not already the last dimension.
      We also don't subtract the max before tf.exp for numerical stability since that might
      mess up the attributions and it seems like TensorFlow doesn't define softmax that way
      (according to the docs)
      """
      in0 = op.inputs[0]
      in0_max = tf.reduce_max(
          in0, axis=-1, keepdims=True, name="in0_max"
      )
      in0_centered = in0 - in0_max
      evals = tf.exp(in0_centered, name="custom_exp")
      rsum = tf.reduce_sum(evals, axis=-1, keepdims=True)
      div = evals / rsum
  
      # mark these as in-between the inputs and outputs
      for op in [evals.op, rsum.op, div.op, in0_centered.op]:
          for t in op.outputs:
              if t.name not in explainer.between_tensors:
                  explainer.between_tensors[t.name] = False
  
      out = tf.gradients(div, in0_centered, grad_ys=grads[0])[0]
  
      # remove the names we just added
      for op in [evals.op, rsum.op, div.op, in0_centered.op]:
          for t in op.outputs:
              if explainer.between_tensors[t.name] is False:
                  del explainer.between_tensors[t.name]
  
      # rescale to account for our shift by in0_max (which we did for numerical stability)
      xin0, rin0 = tf.split(in0, 2)
      xin0_centered, rin0_centered = tf.split(in0_centered, 2)
      delta_in0 = xin0 - rin0
      dup0 = [2] + [1 for i in delta_in0.shape[1:]]
      return tf.where(
          tf.tile(tf.abs(delta_in0), dup0) < 1e-6,
          out,
          out
          * tf.tile((xin0_centered - rin0_centered) / delta_in0, dup0),
      )
  
imports: import warnings
  import numpy
  from packaging import version
  from ...utils._exceptions import DimensionError
  from .._explainer import Explainer
  from ..tf_utils import _get_graph, _get_model_inputs, _get_model_output, _get_session
  from .deep_utils import _check_additivity
label_desc: Use image node_id 5 for calling a global function with example usage: softmax(explainer, op) and returns: tf