code: def brownian_motion_unknown_scales(observed_locs):
      """Brownian Motion model with scale parameters treated as random variables.
  
      Args:
        observed_locs: Array of loc parameters with np.nan value if loc is
          unobserved.
      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;
    }
    parameters {
      real<lower=0> innovation_noise_scale;
      real<lower=0> observation_noise_scale;
      vector[num_timesteps] loc;
    }
    model {
      innovation_noise_scale ~ lognormal(0, 2);
      observation_noise_scale ~ lognormal(0, 2);
      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)],
      }
  
      model = util.cached_stan_model(code)
  
      def _ext_identity(samples):
          """Extracts the values of all latent variables."""
          return {
              "innovation_noise_scale": util.get_columns(
                  samples, r"^innovation_noise_scale$"
              )[:, 0],
              "observation_noise_scale": util.get_columns(
                  samples, r"^observation_noise_scale$"
              )[:, 0],
              "locs": util.get_columns(samples, r"^loc\[\d+\]$"),
          }
  
      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 2 for calling a global function with example usage: brownian_motion_unknown_scales(observed_locs) and returns: stan_model, dict