code: def multinomial_cases():
      num_samples = 10
  
      def wrap(
          op: ops.Operation, logits, logits_dtype, output_dtype, seed
      ):
          device_type = get_device().device_type
          # Some dtypes are not supported on some devices
          if (
              logits_dtype == dtypes.bfloat16
              and device_type == "GPU"
              and not test_util.is_gpu_available(
                  cuda_only=True, min_cuda_compute_capability=(8, 0)
              )
          ):
              logits_dtype = dtypes.float32
          return op(
              seed=seed,
              logits=constant_op.constant(logits, dtype=logits_dtype),
              num_samples=num_samples,
              output_dtype=output_dtype,
          )
  
      for logits_dtype in (
          np.float16,
          dtypes.bfloat16.as_numpy_dtype,
          np.float32,
          np.float64,
      ):
          for output_dtype in dtypes.int32, dtypes.int64:
              for logits in (
                  [[0.1, 0.25, 0.5, 0.15]],
                  [[0.5, 0.5], [0.8, 0.2], [0.25, 0.75]],
              ):
                  yield (
                      "multinomial",
                      functools.partial(
                          wrap,
                          stateless.stateless_multinomial,
                          logits,
                          logits_dtype,
                          output_dtype,
                      ),
                      functools.partial(
                          wrap,
                          random_ops.multinomial,
                          logits,
                          logits_dtype,
                          output_dtype,
                      ),
                  )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.compat import compat
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import random_seed
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import gen_stateless_random_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 5 for calling a global function with example usage: multinomial_cases() and returns: op