code: def _test_dataset(train_size, test_size=None, num_counties=3):
      return dict(
          train_log_uranium=np.random.rand(train_size).astype(
              np.float32
          ),
          train_county=np.arange(train_size) % num_counties,
          train_floor=np.arange(train_size) % 2,
          train_floor_by_county=np.random.rand(train_size).astype(
              np.float32
          ),
          train_log_radon=np.random.rand(train_size).astype(np.float32),
          num_counties=np.array(num_counties, dtype=np.int32),
          test_log_uranium=(
              np.random.rand(test_size).astype(np.float32)
              if test_size
              else None
          ),
          test_county=(
              np.arange(test_size) % num_counties if test_size else None
          ),
          test_floor=(np.arange(test_size) % 2 if test_size else None),
          test_floor_by_county=(
              np.random.rand(test_size).astype(np.float32)
              if test_size
              else None
          ),
          test_log_radon=(
              np.random.rand(test_size).astype(np.float32)
              if test_size
              else None
          ),
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from inference_gym.internal import test_util
  from inference_gym.targets import radon_contextual_effects
label_desc: Use image node_id 1 for calling a global function with example usage: _test_dataset(train_size, test_size, num_counties) and returns: dict