code: def _allowed_types(include_int=False):
      allowed_types = {
          dtypes.float64,
          dtypes.float32,
          dtypes.float16,
          dtypes.bfloat16,
      }
      if include_int:
          allowed_types.update({dtypes.int32, dtypes.int64})
      return allowed_types
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  from tensorflow.compiler.tests import xla_test
  from tensorflow.python.client import device_lib
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.kernel_tests.random import util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_stateless_random_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops_util
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling a global function with example usage: _allowed_types(include_int) and returns: allowed_types