code: def categorical_case(pmf, fns, rand=None):
      """Returns the outputs of fns[i] with probability pmf[i].
  
      Args:
        pmf: A 1-D tensor of probabilities, the probability mass function.
        fns: A list of callables that return tensors, same length as pmf.
        rand: An optional scalar between 0.0 and 1.0, the output of an RNG.
  
      Returns:
        A tensor, the output of fns[i] with probability pmf[i].
      """
      rand = tf.random_uniform([]) if rand is None else rand
      cmf = tf.pad(tf.cumsum(pmf), [(1, 0)])
      cmf = [cmf[i] for i in range(len(fns) + 1)]
      preds = [
          (rand >= a) & (rand < b) for a, b in zip(cmf[:-1], cmf[1:])
      ]
      return tf.case(list(zip(preds, fns)), exclusive=True)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import numpy
  from tensor2tensor.data_generators import problem
  from tensor2tensor.data_generators import text_problems
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: categorical_case(pmf, fns, rand) and returns: tf