code: def to_onehot(i, depth, axis=-1):
      """Convert integer categorical indices to one-hot probability vectors.
  
      Args:
        i: an array of integer categorical indices.
        depth: the number of categories.
        axis: the axis on which to lay out the categories.
  
      Returns:
        An array of one-hot categorical indices, shaped like `i` but with a
        categorical axis in the location specified by `axis`.
      """
      x = np.eye(depth)[i]
      axis %= x.ndim
      if axis != x.ndim - 1:
          # move new axis forward
          axes = list(range(x.ndim - 1))
          axes.insert(axis, x.ndim - 1)
          x = np.transpose(x, axes)
      assert np.allclose(x.sum(axis=axis), 1)
      return x
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import contextlib
  import datetime
  import numbers
  import pdb
  import tempfile
  import time
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling a global function with example usage: to_onehot(i, depth, axis) and returns: x