code: def experiment(
      output_dir: str,
      grid_index: np.ndarray,
      mean_trajectory_length_vals: Optional[np.ndarray] = None,
      damping_vals: Optional[np.ndarray] = None,
  ):
      """Runs an experiment."""
      epath.Path(output_dir).mkdir(parents=True, exist_ok=True)
      if mean_trajectory_length_vals is not None:
          for j, mean_trajectory_length in enumerate(
              mean_trajectory_length_vals
          ):
              logging.info("Starting %d", j)
              whole_grid_index = [grid_index, j]
              res = adaptive_malt.run_grid_element(  # pytype: disable=missing-parameter
                  mean_trajectory_length=mean_trajectory_length,
                  seed=np.random.RandomState(
                      list(whole_grid_index)
                  ).randint(1 << 32),
              )
              utils.save_h5py(
                  os.path.join(
                      output_dir,
                      f"{whole_grid_index[0]}.{whole_grid_index[1]}.h5",
                  ),
                  res,
              )
      elif damping_vals is not None:
          for i, damping in enumerate(damping_vals):
              logging.info("Starting %d", i)
              whole_grid_index = [i, grid_index]
              res = adaptive_malt.run_grid_element(  # pytype: disable=missing-parameter
                  damping=damping,
                  seed=np.random.RandomState(
                      list(whole_grid_index)
                  ).randint(1 << 32),
              )
              utils.save_h5py(
                  os.path.join(
                      output_dir,
                      f"{whole_grid_index[0]}.{whole_grid_index[1]}.h5",
                  ),
                  res,
              )
      else:
          res = adaptive_malt.run_grid_element(  # pytype: disable=missing-parameter
              seed=np.random.RandomState(list(grid_index)).randint(
                  1 << 32
              )
          )
          utils.save_h5py(
              os.path.join(
                  output_dir, f"{grid_index[0]}.{grid_index[1]}.h5"
              ),
              res,
          )
  
file_path: probability/discussion/adaptive_malt/grid_search_runner.py
imports: from collections.abc import Sequence
  import os
  from typing import Optional
  from absl import app
  from absl import flags
  from absl import logging
  from etils import epath
  import gin
  import numpy
  from discussion.adaptive_malt import adaptive_malt
  from discussion.adaptive_malt import utils
label_desc: Use image node_id 1 for calling a global function with example usage: experiment(output_dir, grid_index, mean_trajectory_length_vals, damping_vals) without return types