code: def _swap_axis(logits, dim_index, last_index, **kwargs):
      return tf.transpose(
          logits,
          tf.concat(
              [
                  tf.range(dim_index),
                  [last_index],
                  tf.range(dim_index + 1, last_index),
                  [dim_index],
              ],
              0,
          ),
          **kwargs,
      )
  
imports: import tensorflow
  from tensorflow_addons.utils import types
label_desc: Use image node_id 2 for calling a global function with example usage: _swap_axis(logits, dim_index, last_index) and returns: tf