code: def brownian_motion_log_likelihood_fn(
      values,
      observed_locs,
      use_markov_chain,
      observation_noise_scale=None,
  ):
      """Likelihood of observed data under the Brownian Motion model."""
      if observation_noise_scale is None:
          (_, observation_noise_scale) = values[:2]
          latents = (
              values[2]
              if use_markov_chain
              else tf.stack(values[2:], axis=-1)
          )
      else:
          latents = (
              values if use_markov_chain else tf.stack(values, axis=-1)
          )
  
      observation_noise_scale = tf.convert_to_tensor(
          observation_noise_scale, name="observation_noise_scale"
      )
      is_observed = ~tf.math.is_nan(observed_locs)
      lps = tfd.Normal(
          loc=latents, scale=observation_noise_scale[..., tf.newaxis]
      ).log_prob(tf.where(is_observed, observed_locs, 0.0))
      return tf.reduce_sum(tf.where(is_observed, lps, 0.0), axis=-1)
  
imports: import functools
  import tensorflow.compat.v2
  import tensorflow_probability
  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 brownian_motion_missing_middle_observations
  from inference_gym.targets.ground_truth import brownian_motion_unknown_scales_missing_middle_observations
label_desc: Use image node_id 4 for calling a global function with example usage: brownian_motion_log_likelihood_fn(values, observed_locs, use_markov_chain, observation_noise_scale) and returns: tf