code: def make_radon_prior(num_counties, dtype, prior_scale):
      """Generative process for the radon model with contextual effects."""
      if prior_scale == "uniform":
          county_effect_scale = tfd.Uniform(
              low=tf.zeros([], dtype=dtype), high=100.0
          )
          log_radon_scale = tfd.Uniform(
              low=tf.zeros([], dtype=dtype), high=100.0
          )
      elif prior_scale == "halfnormal":
          county_effect_scale = tfd.HalfNormal(
              scale=tf.ones([], dtype=dtype)
          )
          log_radon_scale = tfd.HalfNormal(
              scale=tf.ones([], dtype=dtype)
          )
      else:
          raise ValueError(
              prior_scale, " is not a valid value for `prior_scale`"
          )
  
      return tfd.JointDistributionNamed(
          dict(
              county_effect_mean=tfd.Normal(
                  loc=tf.zeros([], dtype=dtype), scale=1.0
              ),
              county_effect_scale=county_effect_scale,
              county_effect=(
                  lambda county_effect_scale, county_effect_mean: tfd.Sample(  # pylint: disable=g-long-lambda
                      tfd.Normal(
                          loc=county_effect_mean,
                          scale=county_effect_scale,
                      ),
                      sample_shape=[num_counties],
                  )
              ),
              weight=tfd.Sample(
                  tfd.Normal(loc=tf.zeros([], dtype=dtype), scale=1.0),
                  sample_shape=[3],
              ),
              log_radon_scale=log_radon_scale,
          )
      )
  
imports: import functools
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import dtype_util
  from inference_gym.internal import data
  from inference_gym.targets import bayesian_model
  from inference_gym.targets import model
  from inference_gym.targets.ground_truth import radon_contextual_effects_indiana
  from inference_gym.targets.ground_truth import radon_contextual_effects_indiana_halfnormal
  from inference_gym.targets.ground_truth import radon_contextual_effects_minnesota
  from inference_gym.targets.ground_truth import radon_contextual_effects_minnesota_halfnormal
label_desc: Use image node_id 2 for calling a global function with example usage: make_radon_prior(num_counties, dtype, prior_scale) and returns: tfd