code: def sample(p, axis=None, temperature=1, onehot=False):
      """Sample an array of categorical variables.
  
      Args:
        p: an array of categorical probability vectors, possibly unnormalized.
        axis: the axis that spans the categories (default: -1).
        temperature: if not 1, transform the distribution by dividing the log
            probabilities and renormalizing. Values greater than 1 increase entropy,
            values less than 1 decrease entropy. A value of 0 yields a deterministic
            distribution that chooses the mode.
        onehot: whether to one-hot encode the result.
  
      Returns:
        An array of samples. If `onehot` is False, the result is an array of integer
        category indices, with the categorical axis removed. If `onehot` is True,
        these indices are one-hot encoded, so that the categorical axis remains and
        the result has the same shape and dtype as `p`.
      """
      assert (
          p >= 0
      ).all()  # just making sure we don't put log probabilities in here
  
      if axis is None:
          axis = p.ndim - 1
  
      if temperature != 1:
          p **= 1.0 / temperature
      cmf = p.cumsum(axis=axis)
      totalmasses = cmf[
          tuple(
              slice(None) if d != axis else slice(-1, None)
              for d in range(cmf.ndim)
          )
      ]
      u = np.random.random(
          [p.shape[d] if d != axis else 1 for d in range(p.ndim)]
      )
      i = np.argmax(u * totalmasses < cmf, axis=axis)
  
      return (
          to_onehot(i, axis=axis, depth=p.shape[axis]) if onehot else i
      )
  
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 4 for calling a global function with example usage: sample(p, axis, temperature, onehot) and returns: unknown