code: def categorical_case(pmf, fns, rand=None):
    
    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)
docs: 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].