code: def main(argv):
      if len(argv) > 1:
          raise app.UsageError("Too many command-line arguments.")
  
      tf.enable_v2_behavior()
  
      stan_model = getattr(targets, FLAGS.target)()
  
      with stan_model.sample_fn(
          iter_sampling=FLAGS.stan_samples,
          chains=FLAGS.stan_chains,
          show_progress=True,
      ) as mcmc_output:
          summary = mcmc_output.summary()
          if FLAGS.print_summary:
              pd.set_option("display.max_rows", sys.maxsize)
              pd.set_option("display.max_columns", sys.maxsize)
              print(mcmc_output.diagnose())
              print(summary)
  
          array_strs = []
          for name, fn in sorted(stan_model.extract_fns.items()):
              transformed_samples = []
  
              # We handle one chain at a time to reduce memory usage.
              chain_means = []
              chain_stds = []
              chain_esss = []
              for chain_id in range(FLAGS.stan_chains):
                  # TODO(https://github.com/stan-dev/cmdstanpy/issues/218): This step is
                  # very slow and wastes memory. Consider reading the CSV files ourselves.
  
                  # sample shape is [num_samples, num_chains, num_columns]
                  chain = mcmc_output.draws()[:, chain_id, :]
                  dataframe = pd.DataFrame(
                      chain, columns=mcmc_output.column_names
                  )
  
                  transformed_samples = fn(dataframe)
  
                  # We reduce over the samples dimension. Transformations can return
                  # nested outputs.
                  mean = tf.nest.map_structure(
                      lambda s: s.mean(0), transformed_samples
                  )
                  std = tf.nest.map_structure(
                      lambda s: s.std(0), transformed_samples
                  )
                  ess = tf.nest.map_structure(
                      get_ess, transformed_samples
                  )
  
                  chain_means.append(mean)
                  chain_stds.append(std)
                  chain_esss.append(ess)
  
              # Now we reduce across chains.
              ess = tf.nest.map_structure(
                  lambda *s: np.sum(s, 0), *chain_esss
              )
              mean = tf.nest.map_structure(
                  lambda *s: np.mean(s, 0), *chain_means
              )
              sem = tf.nest.map_structure(
                  lambda std, ess: std / np.sqrt(ess), std, ess
              )
              std = tf.nest.map_structure(
                  lambda *s: np.mean(s, 0), *chain_stds
              )
  
              for (tuple_path, mean_part), sem_part, std_part in zip(
                  nest.flatten_with_tuple_paths(mean),
                  tf.nest.flatten(sem),
                  tf.nest.flatten(std),
              ):
                  array_strs.extend(
                      ground_truth_encoding.save_ground_truth_part(
                          name=name,
                          tuple_path=tuple_path,
                          mean=mean_part,
                          sem=sem_part,
                          std=std_part,
                          sestd=None,
                      )
                  )
  
      argv_str = " \\\n".join(
          ["  {}".format(arg) for arg in sys.argv[1:]]
      )
      command_str = "python -m inference_gym.tools.get_ground_truth \\\n{argv_str}".format(
          argv_str=argv_str
      )
  
      file_str = ground_truth_encoding.get_ground_truth_module_source(
          target_name=FLAGS.target,
          command_str=command_str,
          array_strs=array_strs,
      )
  
      if FLAGS.output_directory is None:
          file_basedir = os.path.dirname(os.path.realpath(__file__))
          output_directory = os.path.join(
              file_basedir, "../targets/ground_truth"
          )
      else:
          output_directory = FLAGS.output_directory
      file_path = os.path.join(
          output_directory, "{}.py".format(FLAGS.target)
      )
      print("Writing ground truth values to: {}".format(file_path))
      with open(file_path, "w") as f:
          f.write(file_str)
  
file_path: probability/spinoffs/inference_gym/inference_gym/tools/get_ground_truth.py
imports: import functools
  import os
  import sys
  from absl import app
  from absl import flags
  import numpy
  import pandas
  import tensorflow.compat.v2
  import tensorflow_probability
  from inference_gym.internal import ground_truth_encoding
  from inference_gym.tools.stan import targets
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: main(argv) without return types