code: def radon(
      state,
      num_examples,
      num_counties,
      drop_floors_above_1=True,
      train_fraction=1.0,
      shuffle=False,
      shuffle_seed=42,
  ):
      """The Radon dataset [1] loaded from `tensorflow_datasets`.
  
      Radon is a radioactive gas that enters homes through contact points with the
      ground. It is a carcinogen that is the primary cause of lung cancer in
      non-smokers. Radon levels vary greatly from household to household.
  
      This dataset contains measured radon concentrations from houses in the
      United States and associated features.
  
      Args:
        state: `str`, the two-character code for the U.S. state by which to filter
          the data. If `None`, data for all states are returned.
        num_examples: `int`, total number of examples in the filtered dataset. When
          the dataset is materialized, this value is verified against the size of
          the filtered dataset, and if there is a mismatch a ValueError is raised.
          (The value is not determined automatically from the dataset because it is
          needed before the dataset is materialized).
        num_counties: `int`, number of unique counties in the filtered dataset. When
          the dataset is materialized, this value is verified against the size of
          the filtered dataset, and if there is a mismatch a ValueError is raised.
          (The value is not determined automatically from the dataset because it is
          needed before the dataset is materialized).
        drop_floors_above_1: Whether to drop datapoints with `floor` being greater
          than 1. In particular, floor equals 9 when it is unknown and the standard
          model uses the floor as a boolean indicator.
        train_fraction: What fraction of the data to put in the training set.
        shuffle: `bool`. If `True`, shuffle the data.
        shuffle_seed: `int`, RNG seed to use if shuffling the data.
  
      Returns:
        dataset: A Dict with the following keys:
          `num_counties`: `int`, number of unique counties in the filtered dataset.
          `train_log_uranium`: Floating-point `Tensor` with shape
            `[num_train]`. Soil uranium measurements.
          `train_floor`: Integer `Tensor` with shape `[num_train]`. Floor of the
            house on which the measurement was taken.
          `train_county`: Integer `Tensor` with values in `range(0, num_counties)`
            of shape `[num_train]`. County in which the measurement was taken.
          `train_floor_by_county`: Floating-point `Tensor` with shape
            `[num_train]`. Average floor on which the measurement was taken for the
            county in which each house is located (the `Tensor` will have
            `num_counties` unique values). This represents the contextual effect.
          `train_log_radon`: Floating-point `Tensor` with shape `[num_train]`.
            Radon measurement for each house (the dependent variable in the model).
          `test_log_uranium`: Floating-point `Tensor` with shape `[num_test`. Soil
            Soil uranium measurements for the test set.
          `test_floor`: Integer `Tensor` with shape `[num_test]`. Floor of the house
            house on which the measurement was taken.
          `test_county`: Integer `Tensor` with values in `range(0, num_counties)` of
            shape `[num_test]`. County in which the measurement was taken. Can be
            `None`, in which case test-related sample transformations are not
            computed.
          `test_floor_by_county`: Floating-point `Tensor` with shape
            `[num_test]`. Average floor on which the measurement was taken
            (calculated from the training set) for the county in which each house is
            located (the `Tensor` will have `num_counties` unique values). This
            represents the contextual effect.
          `test_log_radon`: Floating-point `Tensor` with shape `[num_test]`. Radon
            measurement for each house (the dependent variable in the model).
  
      Raises:
        ValueError if `num_examples` is not equal to the number of examples in the
          materialized dataset.
        ValueError if `num_counties` is not equal to the number of unique counties
          in the materialized dataset.
  
      #### References
  
      [1] Gelman, A., & Hill, J. (2007). Data Analysis Using Regression and
          Multilevel/Hierarchical Models (1st ed.). Cambridge University Press.
          http://www.stat.columbia.edu/~gelman/arm/examples/radon/
      """
      num_train = int(num_examples * train_fraction)
      num_test = num_examples - num_train
  
      def load_dataset():
          """Function that actually loads the dataset."""
          if load_dataset.dataset is not None:
              return load_dataset.dataset
  
          with tf.name_scope("radon"), tf.init_scope():
              dataset = _tfds().load(
                  name="radon:1.*.*", split="train", batch_size=-1
              )
              dataset = _tfds().as_numpy(dataset)
  
              states = dataset["features"]["state"].astype("U13")
              floor = dataset["features"]["floor"].astype(np.int32)
              radon_val = dataset["activity"].astype(np.float64)
              county_strings = dataset["features"]["county"].astype(
                  "U13"
              )
              uranium = dataset["features"]["Uppm"].astype(np.float64)
  
              if state is not None:
                  floor = floor[states == state]
                  radon_val = radon_val[states == state]
                  county_strings = county_strings[states == state]
                  uranium = uranium[states == state]
  
              if drop_floors_above_1:
                  keep = floor <= 1
                  floor = floor[keep]
                  radon_val = radon_val[keep]
                  county_strings = county_strings[keep]
                  uranium = uranium[keep]
  
              radon_val[radon_val <= 0.0] = 0.1
              log_radon = np.log(radon_val)
              log_uranium = np.log(uranium)
              unique_counties, county = np.unique(
                  county_strings, return_inverse=True
              )
              county = county.astype(np.int32)
  
              if log_radon.size != num_examples:
                  raise ValueError(
                      "The size of the filtered dataset must equal the input "
                      "`num_examples`. Saw dataset size = {}, `num_examples` = {}"
                      "".format(log_radon.size, num_examples)
                  )
              if unique_counties.size != num_counties:
                  raise ValueError(
                      "The number of counties present in the filtered dataset must equal "
                      "the input `num_counties`. Saw {} counties but `num_counties` = {}"
                      "".format(unique_counties.size, num_counties)
                  )
  
              if shuffle:
                  shuffle_idxs = np.arange(num_examples)
                  np.random.RandomState(shuffle_seed).shuffle(
                      shuffle_idxs
                  )
                  log_uranium = log_uranium[shuffle_idxs]
                  floor = floor[shuffle_idxs]
                  county = county[shuffle_idxs]
                  log_radon = log_radon[shuffle_idxs]
  
              train_floor = floor[:num_train]
              train_county = county[:num_train]
              test_floor = floor[num_train:]
              test_county = county[num_train:]
  
              # Create a new features for mean of floor across counties.
              xbar = []
              for i in range(num_counties):
                  xbar.append(floor[county == i].mean())
              floor_by_county = np.array(xbar, dtype=log_radon.dtype)
  
              load_dataset.dataset = dict(
                  train_log_uranium=log_uranium[:num_train],
                  train_floor=train_floor,
                  train_county=train_county,
                  train_floor_by_county=floor_by_county[train_county],
                  train_log_radon=log_radon[:num_train],
                  test_log_uranium=log_uranium[num_train:],
                  test_floor=test_floor,
                  test_county=test_county,
                  test_floor_by_county=floor_by_county[test_county],
                  test_log_radon=log_radon[num_train:],
              )
  
          return load_dataset.dataset
  
      load_dataset.dataset = None
  
      return dict(
          num_counties=np.array(num_counties, dtype=np.int32),
          train_log_uranium=_defer(
              lambda: load_dataset()["train_log_uranium"],
              shape=[num_train],
              dtype=np.float64,
          ),
          train_floor=_defer(
              lambda: load_dataset()["train_floor"],
              shape=[num_train],
              dtype=np.int32,
          ),
          train_county=_defer(
              lambda: load_dataset()["train_county"],
              shape=[num_train],
              dtype=np.int32,
          ),
          train_floor_by_county=_defer(
              lambda: load_dataset()["train_floor_by_county"],
              shape=[num_train],
              dtype=np.float64,
          ),
          train_log_radon=_defer(
              lambda: load_dataset()["train_log_radon"],
              shape=[num_train],
              dtype=np.float64,
          ),
          test_log_uranium=_defer(
              lambda: load_dataset()["test_log_uranium"],
              shape=[num_test],
              dtype=np.float64,
          ),
          test_floor=_defer(
              lambda: load_dataset()["test_floor"],
              shape=[num_test],
              dtype=np.int32,
          ),
          test_county=_defer(
              lambda: load_dataset()["test_county"],
              shape=[num_test],
              dtype=np.int32,
          ),
          test_floor_by_county=_defer(
              lambda: load_dataset()["test_floor_by_county"],
              shape=[num_test],
              dtype=np.float64,
          ),
          test_log_radon=_defer(
              lambda: load_dataset()["test_log_radon"],
              shape=[num_test],
              dtype=np.float64,
          ),
      )
  
imports: import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from inference_gym.internal.datasets import brownian_motion_missing_middle_observations
  from inference_gym.internal.datasets import convection_lorenz_bridge
  from inference_gym.internal.datasets import sp500_closing_prices
  from inference_gym.internal.datasets import synthetic_item_response_theory
  from inference_gym.internal.datasets import synthetic_log_gaussian_cox_process
  from inference_gym.internal.datasets import synthetic_plasma_spectroscopy
  from inference_gym.internal.datasets import synthetic_plasma_spectroscopy_with_bump
label_desc: Use image node_id 7 for calling a global function with example usage: radon(state, num_examples, num_counties, drop_floors_above_1, train_fraction, shuffle, shuffle_seed) and returns: dict, load_dataset, load_dataset