code: def sample(logits):
      """
      Creates a sampling Tensor for non deterministic policies
      when using categorical distribution.
      It uses the Gumbel-max trick: http://amid.fish/humble-gumbel
  
      :param logits: (TensorFlow Tensor) The input probability for each action
      :return: (TensorFlow Tensor) The sampled action
      """
      noise = tf.random_uniform(tf.shape(logits))
      return tf.argmax(logits - tf.log(-tf.log(noise)), 1)
  
imports: import os
  import collections
  import functools
  import multiprocessing
  from typing import Set
  import numpy
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: sample(logits) and returns: tf