code: def softmax(explainer, op, *grads):
      
      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
      explainer.between_ops.extend([evals.op, rsum.op, div.op, in0_centered.op]) # mark these as in-between the inputs and outputs
      out = tf.gradients(div, in0_centered, grad_ys=grads[0])[0]
      del explainer.between_ops[-4:]
  
      # 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)
      )
docs: 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)