code: def brownian_motion(
      observed_locs, innovation_noise_scale, observation_noise_scale
  ):
      """Brownian Motion model.
  
      Args:
        observed_locs: Array of loc parameters with np.nan value if loc is
          unobserved.
        innovation_noise_scale: Python `float`.
        observation_noise_scale: Python `float`.
  
      Returns:
        model: `StanModel`.
      """
  
      code = """
    data {
      int<lower=0> num_timesteps;
      int<lower=0> num_observations;
      int<lower = 1, upper = num_timesteps> observation_indices[num_observations];
      vector[num_observations] observations;
      real<lower=0> innovation_noise_scale;
      real<lower=0> observation_noise_scale;
    }
    parameters {
      vector[num_timesteps] loc;
    }
    model {
      loc[1] ~ normal(0, innovation_noise_scale);
      for (t in 2:num_timesteps){
        loc[t] ~ normal(loc[t-1], innovation_noise_scale);
        }
      observations ~ normal(loc[observation_indices], observation_noise_scale);
    }
    """
  
      stan_data = {
          "num_timesteps": len(observed_locs),
          "num_observations": len(
              observed_locs[~np.isnan(observed_locs)]
          ),
          "observation_indices": (
              np.arange(1, len(observed_locs) + 1)[
                  ~np.isnan(observed_locs)
              ]
          ),
          "observations": observed_locs[~np.isnan(observed_locs)],
          "innovation_noise_scale": innovation_noise_scale,
          "observation_noise_scale": observation_noise_scale,
      }
  
      model = util.cached_stan_model(code)
  
      def _ext_identity(samples):
          """Extracts the values of all latent variables."""
          locs = util.get_columns(samples, r"^loc\[\d+\]$")
          return locs
  
      extract_fns = {"identity": _ext_identity}
  
      return stan_model.StanModel(
          extract_fns=extract_fns,
          sample_fn=util.make_sample_fn(model, data=stan_data),
      )
  
imports: import numpy
  from inference_gym.tools.stan import stan_model
  from inference_gym.tools.stan import util
label_desc: Use image node_id 1 for calling a global function with example usage: brownian_motion(observed_locs, innovation_noise_scale, observation_noise_scale) and returns: stan_model, locs