code: def softmax(x, axis=-1):
      x = x - tf.reduce_max(input_tensor=x, axis=axis, keepdims=True)
      ex = tf.exp(x)
      return ex / tf.reduce_sum(
          input_tensor=ex, axis=axis, keepdims=True
      )
  
file_path: finetune/finetune/base_models/gpt2/featurizer.py
imports: import functools
  import tensorflow
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.activations import gelu
  from finetune.base_models.gpt.featurizer import norm, dropout, get_pos_values
label_desc: Use image node_id 1 for calling a global function with example usage: softmax(x, axis) and returns: unknown