code: def run_ax_search(
      fixed_params: Dict,
      ax_params: List[Dict[str, Any]],
      eval_fn: Callable,
      obj_name: str,
      minimize: bool,
      id_: str,
      parse_params_fn: Optional[Callable] = None,
      ax_param_constraints: Optional[List[str]] = None,
      num_ax_steps: int = 50,
      num_concur_samples: int = 2,
      num_seeds: int = 10,
      num_proc: int = 20,
      folder_name: Optional[str] = None,
      verbose: bool = False,
  ) -> Tuple[Dict[str, Any], AxClient]:
      """
      Run a search for best hyperparameter values using Ax.
      Note that this requires the Ax package (https://ax.dev/) to be installed.
  
      Args:
          fixed_params: Fixed values of hyperparameters.
          ax_params: Ax configuration for hyperparameters that are searched over. See docs for ax_client.create_experiment()
          eval_fn: Evaluation function that returns a dictionary of metric values.
          obj_name: Objective name (key of the dict returned by eval_fn)
          minimize: If True, objective is minimized, if False it's maximized.
          id_: An arbitrary string identifier of the search (used as part of filename where results are saved)
          parse_params_fn: A function applied to the parameter dictionary to parse it. Can be used
              if the best representation for Ax doesn't match the format accepted by the eval_fn.
          ax_param_constraints: Constraints for the parameters that are searched over.
          num_ax_steps: The number of ax steps to take.
          num_concur_samples: Number of configurations to sample per ax step (in parallel)
          num_seeds: Number of seeds to average over
          num_proc: Number of processes to run in parallel.
          folder_name: Folder where to save best found parameters
          verbose: If True, some details are printed out
      Returns:
          A dict of best hyperparameters found by Ax
      """
      for p in ax_params:
          assert (
              p["name"] not in fixed_params
          ), f'Parameter {p["name"]} appers in both fixed and search parameters'
      if ax_param_constraints is None:
          ax_param_constraints = []
      ax_client = AxClient()
      ax_client.create_experiment(
          name=f"hparams_search_{id_}",
          parameters=ax_params,
          objective_name=obj_name,
          minimize=minimize,
          parameter_constraints=ax_param_constraints,
          choose_generation_strategy_kwargs={
              "max_parallelism_override": num_concur_samples,
              "num_initialization_trials": max(
                  num_concur_samples, 5, len(ax_params)
              ),
          },
      )
      best_params = None
      all_considered_params = []
      all_considered_metrics = []
  
      try:
          for i in range(1, num_ax_steps + 1):
              if verbose:
                  print(f"ax step {i}/{num_ax_steps}")
              params_list = []
              trial_indices_list = []
              for _ in range(num_concur_samples):
                  # sample several values (to be evaluated in parallel)
                  parameters, trial_index = ax_client.get_next_trial()
                  params_list.append(parameters)
                  trial_indices_list.append(trial_index)
              res = ax_evaluate_params(
                  params_list,
                  fixed_params=fixed_params,
                  eval_fn=eval_fn,
                  parse_params_fn=parse_params_fn,
                  num_seeds=num_seeds,
                  num_proc=num_proc,
              )
              all_considered_params.extend(params_list)
              all_considered_metrics.extend(res)
              for t_i, v in zip(trial_indices_list, res):
                  ax_client.complete_trial(trial_index=t_i, raw_data=v)
              (
                  best_params,
                  predicted_metrics,
              ) = ax_client.get_best_parameters()
              predicted_metrics = predicted_metrics[
                  0
              ]  # choose expected metric values
              if verbose:
                  print(best_params, predicted_metrics)
              # save at every iteration in case search is interrupted
              if folder_name is not None:
                  with open(
                      os.path.join(
                          os.path.expanduser(folder_name),
                          f"ax_results_{id_}.json",
                      ),
                      "w",
                  ) as f:
                      json.dump(
                          {
                              "best_params": best_params,
                              "predicted_metrics": predicted_metrics,
                              "fixed_params": fixed_params,
                              "ax_params": ax_params,
                              "num_ax_steps": i,
                              "num_concur_samples": num_concur_samples,
                              "num_seeds": num_seeds,
                              "num_proc": num_proc,
                              "all_considered_params": all_considered_params,
                              "all_considered_metrics": all_considered_metrics,
                          },
                          f,
                          indent=4,
                      )
      except KeyboardInterrupt:
          # handle keyboard interruption to enable returning intermediate results if interrupted
          pass
      return best_params, ax_client
  
file_path: ReAgent/reagent/scripts/hparam_tuning.py
imports: import logging
  import copy
  import json
  import os
  from typing import Any, Callable, Dict, List, Optional, Tuple
  import numpy
  import torch.multiprocessing
  from ax.service.ax_client import AxClient
label_desc: Use image node_id 2 for calling a global function with example usage: run_ax_search(fixed_params, ax_params, eval_fn, obj_name, minimize, id_, parse_params_fn, ax_param_constraints, num_ax_steps, num_concur_samples, num_seeds, num_proc, folder_name, verbose) and returns: best_params,  ax_client