code: def test_get_suggestion_beam_size_gt_one(self, get_architecture):
      # Trials should be ranked by trial_id. That is, we should only ever fork
      # from the first 2 trials.
      beam_size = 2
      trial_to_arch = {
          0: np.array([1, 2, 3]),
          1: np.array([4, 5, 6]),
          2: np.array([7, 8, 9]),
          3: np.array([10, 11, 12]),
      }
      get_architecture.side_effect = lambda idx: trial_to_arch[int(idx)]
      spec = search_test_utils.create_spec(
          phoenix_spec_pb2.PhoenixSpec.DNN,
          blocks_to_use=["FULLY_CONNECTED_PYRAMID"],
      )
      spec.beam_size = beam_size
      spec.increase_complexity_minimum_trials.append(0)
      algorithm = coordinate_descent.CoordinateDescent(
          spec, self._metadata
      )
      hparams = hp.HParams(
          new_block_type="FIXED_CHANNEL_CONVOLUTION_16"
      )  # enum: 1
  
      # Create one fake trial for each architecture.
      trials = []
      for i in range(3):
          trials.append(
              trial_module.Trial(
                  {
                      "id": i,
                      "model_dir": str(i),
                      "status": "COMPLETED",
                      "trial_infeasible": False,
                      "final_measurement": {"objective_value": i},
                  }
              )
          )
  
      # Adding one infeasible to make sure we don't fork from it.
      trials.append(
          trial_module.Trial(
              {
                  "id": 4,
                  "model_dir": "4",
                  "status": "COMPLETED",
                  "trial_infeasible": True,
                  "final_measurement": {"objective_value": 0},
              }
          )
      )
      logging.info(trials)
  
      # Since forking is random, fork 1000 times then check that we forked from
      # only the trials we care about.
      forked = set()
      for i in range(1000):
          _, fork_trial = algorithm.get_suggestion(trials, hparams)
          forked.add(int(fork_trial))
      self.assertEqual(forked, {0, 1})
  
file_path: model_search/model_search/search/coordinate_descent_test.py
imports: from absl import logging
  from absl.testing import parameterized
  import mock
  from model_search import hparam
  from model_search.metadata import ml_metadata_db
  from model_search.metadata import trial
  from model_search.proto import phoenix_spec_pb2
  from model_search.search import coordinate_descent
  from model_search.search import test_utils
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 3 for calling the CoordinateDescentTest obj's underlying member method code with example usage: obj.test_get_suggestion_beam_size_gt_one(get_architecture) without return types