code: def make_radon_observation_dist(
      params, log_uranium, floor, county, floor_by_county
  ):
      """Likelihood of observed data under the contextual effects radon model."""
      floor = tf.cast(floor, dtype=log_uranium.dtype)
      return tfd.Normal(
          loc=affine(
              log_uranium,
              params["weight"][..., :1],
              affine(floor, params["weight"][..., 1:2])
              + affine(floor_by_county, params["weight"][..., 2:])
              + tf.gather(params["county_effect"], county, axis=-1),
          ),
          scale=params["log_radon_scale"][..., tf.newaxis],
      )
  
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 3 for calling a global function with example usage: make_radon_observation_dist(params, log_uranium, floor, county, floor_by_county) and returns: tfd