code: def softmax_common(x, axis, use_fast_exp):
      """The common part of softmax and fast_softmax"""
      shape = x.shape
      if axis < 0:
          axis = len(shape) + axis
      if axis >= len(shape):
          ValueError("axis parameter should be less than input dim")
  
      k1 = te.reduce_axis((0, shape[axis]), name="k")
      k2 = te.reduce_axis((0, shape[axis]), name="k")
  
      def insert_reduce_index(indices, reduce_index):
          return indices[:axis] + (reduce_index,) + indices[axis:]
  
      def get_non_reduce_indices(indices):
          return tuple(
              [var for (i, var) in enumerate(indices) if i != axis]
          )
  
      def _compute_max(*indices):
          eval_range = insert_reduce_index(indices, k1)
          return tvm.te.max(x[eval_range], axis=k1)
  
      def _compute_delta(max_elem, *indices):
          non_reduce_indices = get_non_reduce_indices(indices)
          return x[indices] - max_elem[non_reduce_indices]
  
      def _compute_exp(max_elem, *indices):
          non_reduce_indices = get_non_reduce_indices(indices)
          return te.exp(x[indices] - max_elem[non_reduce_indices])
  
      def _compute_expsum(exp, *indices):
          eval_range = insert_reduce_index(indices, k2)
          return te.sum(exp[eval_range], axis=k2)
  
      def _normalize(exp, expsum, *indices):
          non_reduce_indices = get_non_reduce_indices(indices)
          return exp[indices] / expsum[non_reduce_indices]
  
      reduced_shape = tuple(
          [dim for (i, dim) in enumerate(shape) if i != axis]
      )
      max_elem = te.compute(
          reduced_shape, _compute_max, name="T_softmax_maxelem"
      )
  
      if use_fast_exp:
          delta = te.compute(
              shape,
              lambda *indices: _compute_delta(max_elem, *indices),
              name="T_softmax_delta",
          )
          exp = topi.math.fast_exp(delta)
      else:
          exp = te.compute(
              shape,
              lambda *indices: _compute_exp(max_elem, *indices),
              name="T_softmax_exp",
          )
      expsum = te.compute(
          reduced_shape,
          lambda *indices: _compute_expsum(exp, *indices),
          name="T_softmax_expsum",
      )
      return te.compute(
          shape,
          lambda *indices: _normalize(exp, expsum, *indices),
          name="T_softmax_norm",
          attrs={"axis": axis},
      )
  
imports: from __future__ import absolute_import
  import tvm
  from tvm import te, topi
label_desc: Use image node_id 3 for calling a global function with example usage: softmax_common(x, axis, use_fast_exp) and returns: te, unknown, tuple, tvm, unknown, te, te, unknown