docs: def config_generator(search_space, max_search, rng, shuffle=True):
          """Generates config dicts from the given search space
  
          Args:
              search_space: (dict) A dictionary of parameters to search over.
                  See note below for more details.
              max_search: (int) The maximum number of configurations to search.
                  If max_search is None, do a full grid search of all discrete
                      parameters, filling in range parameters as needed.
                  Otherwise, do a full grid search of all discrete
                      parameters and then cycle through again filling in new
                      range parameters values; if there are no range parameters,
                      stop after yielding the full cross product of parameters
                      once.
              shuffle: (bool) If True, shuffle the order of generated configs
  
          Yields:
              configs: each config is a dict of parameter values based on the
                  provided search space
  
          The search_space dictionary may consist of two types of parameters:
              --discrete: a discrete parameter is either a single value or a
                  list of values. Use single values, for example, to override
                  a default model parameter or set a flag such as 'verbose'=True.
              --range: a range parameter is a dict of the form:
                  {'range': [<min>, <max>], 'scale': <scale>}
                  where <min> and <max> are the min/max values to search between
                  and scale is one of ['linear', 'log'] (defaulting to 'linear')
                  representing the scale to use when searching the given range
  
          Example:
              search_space = {
                  'verbose': True,                              # discrete
                  'n_epochs': 100,                              # discrete
                  'momentum': [0.0, 0.9, 0.99],                       # discrete
                  'l2': {'range': [0.0001, 10]}                 # linear range
                  'lr': {'range': [0.001, 1], 'scale': 'log'},  # log range
              }
              If max_search is None, this will return 3 configurations (enough to
                  just cover the full cross-product of discrete values, filled
                  in with sampled range values)
              Otherewise, this will return max_search configurations
                  (cycling through the discrete value combinations multiple
                  time if necessary)
          """
  
          def dict_product(d):
              keys = d.keys()
              for element in product(*d.values()):
                  yield dict(zip(keys, element))
  
          def range_param_func(v):
              scale = v.get("scale", "linear")
              mini = min(v["range"])
              maxi = max(v["range"])
              if scale == "linear":
                  func = lambda rand: mini + (maxi - mini) * rand
              elif scale == "log":
                  mini = np.log(mini)
                  maxi = np.log(maxi)
                  func = lambda rand: np.exp(mini + (maxi - mini) * rand)
              else:
                  raise ValueError(f"Unrecognized scale '{scale}' for " "parameter {k}")
              return func
  
          discretes = {}
          ranges = {}
          for k, v in search_space.items():
              if isinstance(v, dict):
                  ranges[k] = range_param_func(v)
              elif isinstance(v, list):
                  discretes[k] = v
              else:
                  discretes[k] = [v]
  
          discrete_configs = list(dict_product(discretes))
  
          if shuffle:
              rng.shuffle(discrete_configs)
  
          # If there are range parameters and a non-None max_search, cycle
          # through the discrete_configs (with new range values) until
          # max_search is met
          if ranges and max_search:
              discrete_configs = cycle(discrete_configs)
  
          for i, config in enumerate(discrete_configs):
              # We may see the same config twice due to cycle
              config = config.copy()
              if max_search and i == max_search:
                  break
              for k, v in ranges.items():
                  config[k] = float(v(rng.random()))
              yield config