class_docstr: Construct the Brownian Motion model.
  
  This models a Brownian Motion process with a Gaussian observation model.
  
  ```none
  locs[0] ~ Normal(loc=0, scale=innovation_noise_scale)
  for t in range(1, num_timesteps):
    locs[t] ~ Normal(loc=locs[t - 1], scale=innovation_noise_scale)
  
  for t in range(num_timesteps):
    observed_locs[t] ~ Normal(loc=locs[t], scale=observation_noise_scale)
  ```
  
  This model supports missing observations, indicated by NaNs in the
  `observed_locs` parameter.
class_docstr_tok: ["Construct","the","Brownian","Motion","model",".","This","models","a","Brownian","Motion","process","with","a","Gaussian","observation","model",".","``","`","none","locs","[","0","]","~","Normal","(","loc=0",",","scale=innovation_noise_scale",")","for","t","in","range","(","1",",","num_timesteps",")",":","locs","[","t","]","~","Normal","(","loc=locs","[","t","-","1","]",",","scale=innovation_noise_scale",")","for","t","in","range","(","num_timesteps",")",":","observed_locs","[","t","]","~","Normal","(","loc=locs","[","t","]",",","scale=observation_noise_scale",")","``","`","This","model","supports","missing","observations",",","indicated","by","NaNs","in","the","`","observed_locs","`","parameter","."]
docstr: Construct the Brownian Motion model.
  
  Args:
    observed_locs: Array of loc parameters with NaN value if loc is
      unobserved.
    innovation_noise_scale: Python `float`.
    observation_noise_scale: Python `float`.
    use_markov_chain: Python `bool` indicating whether to use the
      `MarkovChain` distribution in place of separate random variables for
      each time step. The default of `False` is for backwards compatibility;
      setting this to `True` should significantly improve performance.
    name: Python `str` name prefixed to Ops created by this class.
    pretty_name: A Python `str`. The pretty name of this model.
docstr_tok: ["Construct","the","Brownian","Motion","model",".","Args",":","observed_locs",":","Array","of","loc","parameters","with","NaN","value","if","loc","is","unobserved",".","innovation_noise_scale",":","Python","`","float","`",".","observation_noise_scale",":","Python","`","float","`",".","use_markov_chain",":","Python","`","bool","`","indicating","whether","to","use","the","`","MarkovChain","`","distribution","in","place","of","separate","random","variables","for","each","time","step",".","The","default","of","`","False","`","is","for","backwards","compatibility",";","setting","this","to","`","True","`","should","significantly","improve","performance",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","pretty_name",":","A","Python","`","str","`",".","The","pretty","name","of","this","model","."]
code: def __init__(
      self,
      observed_locs,
      innovation_noise_scale,
      observation_noise_scale,
      use_markov_chain=False,
      name="brownian_motion",
      pretty_name="Brownian Motion",
  ):
      """Construct the Brownian Motion model.
  
      Args:
        observed_locs: Array of loc parameters with NaN value if loc is
          unobserved.
        innovation_noise_scale: Python `float`.
        observation_noise_scale: Python `float`.
        use_markov_chain: Python `bool` indicating whether to use the
          `MarkovChain` distribution in place of separate random variables for
          each time step. The default of `False` is for backwards compatibility;
          setting this to `True` should significantly improve performance.
        name: Python `str` name prefixed to Ops created by this class.
        pretty_name: A Python `str`. The pretty name of this model.
      """
      with tf.name_scope(name):
          num_timesteps = observed_locs.shape[0]
          if use_markov_chain:
              self._prior_dist = brownian_motion_as_markov_chain(
                  num_timesteps=num_timesteps,
                  innovation_noise_scale=innovation_noise_scale,
              )
          else:
              self._prior_dist = tfd.JointDistributionCoroutine(
                  functools.partial(
                      brownian_motion_prior_fn,
                      num_timesteps=num_timesteps,
                      innovation_noise_scale=innovation_noise_scale,
                  )
              )
  
          self._log_likelihood_fn = functools.partial(
              brownian_motion_log_likelihood_fn,
              observation_noise_scale=observation_noise_scale,
              observed_locs=observed_locs,
              use_markov_chain=use_markov_chain,
          )
  
          def _ext_identity(params):
              return tf.stack(params, axis=-1)
  
          def _ext_identity_markov_chain(params):
              return params
  
          sample_transformations = {
              "identity": model.Model.SampleTransformation(
                  fn=(
                      _ext_identity_markov_chain
                      if use_markov_chain
                      else _ext_identity
                  ),
                  pretty_name="Identity",
              )
          }
  
      if use_markov_chain:
          event_space_bijector = tfb.Identity()
      else:
          event_space_bijector = type(self._prior_dist.dtype)(
              *([tfb.Identity()] * num_timesteps)
          )
      super(BrownianMotion, self).__init__(
          default_event_space_bijector=event_space_bijector,
          event_shape=self._prior_dist.event_shape,
          dtype=self._prior_dist.dtype,
          name=name,
          pretty_name=pretty_name,
          sample_transformations=sample_transformations,
      )
  
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 1 to create a new BrownianMotion object from inherited base classes: bayesian_model with example: obj = BrownianMotion(observed_locs, innovation_noise_scale, observation_noise_scale, use_markov_chain, name, pretty_name)