class_docstr: Phoenix: A smart search AutoML algorithm.
class_docstr_tok: ["Phoenix",":","A","smart","search","AutoML","algorithm","."]
docstr: Returns a model_fn for the estimator.
docstr_tok: ["Returns","a","model_fn","for","the","estimator","."]
code: def _make_model_fn(self, run_config, train_steps, use_tpu=False):
      """Returns a model_fn for the estimator."""
  
      def model_fn(features, labels, mode, params):
          """Model function that wraps the model specified."""
          self._metric_fn = self._user_specified_metric_fn
          self._default_metric_fn_list = []
          if self._logits_dimension >= 2:
              self._default_metric_fn_list.append(
                  metric_fns.make_accuracy_metric_fn(
                      self._label_vocabulary
                  )
              )
          if self._logits_dimension == 2:
              self._default_metric_fn_list += [
                  metric_fns.make_auc_roc_metric_fn(
                      self._label_vocabulary
                  ),
                  metric_fns.make_auc_pr_metric_fn(
                      self._label_vocabulary
                  ),
              ]
  
          my_id = architecture_utils.DirectoryHandler.get_trial_id(
              run_config.model_dir, self._phoenix_spec
          )
  
          # Create a copy of hparams
          hparams = params
          if my_id <= len(self._phoenix_spec.user_suggestions):
              hparams = _merge_hparams(
                  params,
                  hp.HParams.from_proto(
                      self._phoenix_spec.user_suggestions[
                          my_id - 1
                      ].hparams
                  ),
              )
  
          # When predicting for RNN, we might not need the length.
          is_training = mode == tf.estimator.ModeKeys.TRAIN
          lengths_feature_name = self._phoenix_spec.lengths_feature_name
          if mode == tf.estimator.ModeKeys.PREDICT:
              if (
                  isinstance(features, dict)
                  and lengths_feature_name not in features
              ):
                  lengths_feature_name = ""
  
          if (
              "params"
              in inspect.signature(self._input_layer_fn).parameters
          ):
              input_layer_fn = functools.partial(
                  self._input_layer_fn,
                  features=features,
                  params=hparams,
                  lengths_feature_name=lengths_feature_name,
              )
          else:
              input_layer_fn = functools.partial(
                  self._input_layer_fn,
                  features=features,
                  lengths_feature_name=lengths_feature_name,
              )
  
          towers, trials = self.keras_model_builder(
              hparams=hparams,
              is_training=is_training,
              run_config=run_config,
              input_layer_fn=input_layer_fn,
              compile_model=False,
          )
  
          for generator_towers in towers.values():
              for tower in generator_towers:
                  tower(None, training=is_training)
  
          trial_mode = trial_utils.get_trial_mode(
              self._phoenix_spec.ensemble_spec,
              self._phoenix_spec.distillation_spec,
              my_id,
          )
  
          training_hooks = []
          # TODO(b/172564129): Figure out how to handle transfer learning for multi
          # task. Install transfer learning hook.
          if (
              is_training
              and self._phoenix_spec.transfer_learning_spec.transfer_learning_type
              in _TL_HOOKS
          ):
              tower_name = base_tower_generator.SEARCH_GENERATOR
              vars_to_warm_start = (
                  architecture_utils.get_tower_variables(tower_name)
              )
              if vars_to_warm_start:
                  hook_fn = _TL_HOOKS[
                      self._phoenix_spec.transfer_learning_spec.transfer_learning_type
                  ]
                  tl_spec = self._phoenix_spec.transfer_learning_spec
                  tl_hook = hook_fn(
                      vars_to_warm_start=vars_to_warm_start,
                      current_trial_id=my_id,
                      completed_trials=trials,
                      discount_factor=tl_spec.previous_trials_discount_factor,
                      max_completed_trials=tl_spec.max_completed_trials,
                      model_dir=run_config.model_dir,
                  )
                  training_hooks.append(tl_hook)
  
          learning_rate_spec_keys = [
              "learning_rate",
              "l2_regularization",
              "gradient_max_norm",
              "exponential_decay_steps",
              "exponential_decay_rate",
          ]
          learning_rate_spec = {
              key: value
              for key, value in hparams.values().items()
              if key in learning_rate_spec_keys
          }
          tower_name = None
  
          # Create the metric_fn if it wasn't specified.
          if not self._metric_fn:
              metric_fn = metric_fns.create_num_parameters_metric_fn(
                  tower_name
              )
              self._default_metric_fn_list.append(metric_fn)
              self._metric_fn = metric_fns.combine_metric_fns(
                  self._default_metric_fn_list
              )
  
          model_spec = self._task_manager.create_model_spec(
              features=features,
              params=hparams,
              learning_rate_spec=learning_rate_spec,
              use_tpu=use_tpu,
              trial_mode=trial_mode,
              towers=towers,
              labels=labels,
              mode=mode,
              my_id=my_id,
              model_directory=run_config.model_dir,
              predictions_fn=self._predictions_fn,
              optimizer_fn=self._user_optimizer_fn,
          )
  
          if run_config.is_chief:
              self._metadata.after_generating_trial_model(my_id)
              search_architecture = [["no_search"]]
              if base_tower_generator.SEARCH_GENERATOR in towers.keys():
                  search_architecture = [
                      t.architecture
                      for t in towers[
                          base_tower_generator.SEARCH_GENERATOR
                      ]
                  ]
              trial_utils.write_replay_spec(
                  model_dir=run_config.model_dir,
                  filename=REPLAY_CONFIG_FILENAME,
                  original_spec=self._phoenix_spec,
                  search_architecture=search_architecture[0],
                  hparams=hparams,
              )
  
          # No need to add train op for the eval graph.
          train_op = None
          if is_training:
              train_op = self._increment_global_step(
                  model_spec.train_op,
                  train_steps,
                  base_tower_generator.SEARCH_GENERATOR,
              )
  
          if (
              isinstance(labels, dict)
              and not self._phoenix_spec.pass_label_dict_as_is
          ):
              label_names = [
                  label_spec.label_name
                  for label_spec in self._phoenix_spec.multi_task_spec
              ]
              actual_labels = {
                  name: label
                  for name, label in labels.items()
                  if name in label_names
              }
          else:
              actual_labels = labels
  
          if use_tpu:
              eval_metrics = None
              weights = None
              if self._phoenix_spec.weight_feature_name:
                  weights = features[
                      self._phoenix_spec.weight_feature_name
                  ]
              if (
                  mode != tf.estimator.ModeKeys.PREDICT
                  and not self._head
              ):
                  eval_metrics = (
                      self._metric_fn,
                      [actual_labels, model_spec.predictions, weights],
                  )
              return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
                  mode=mode,
                  loss=model_spec.loss,
                  predictions=model_spec.predictions,
                  train_op=train_op,
                  eval_metrics=eval_metrics,
                  training_hooks=training_hooks,
              )
  
          eval_metric_ops = None
          if mode != tf.estimator.ModeKeys.PREDICT and not self._head:
              weights = None
              if self._phoenix_spec.weight_feature_name:
                  weights = features[
                      self._phoenix_spec.weight_feature_name
                  ]
              eval_metric_ops = self._metric_fn(
                  actual_labels, model_spec.predictions, weights
              )
          if self._head:
              return self._head.create_estimator_spec(
                  features,
                  mode,
                  model_spec.eval_logits,
                  labels,
                  train_op_fn=lambda _: train_op,
              )
  
          return tf.estimator.EstimatorSpec(
              mode=mode,
              loss=model_spec.loss,
              predictions=model_spec.predictions,
              train_op=train_op,
              training_hooks=training_hooks + model_spec.train_hooks,
              eval_metric_ops=eval_metric_ops,
          )
  
      return model_fn
  
imports: import functools
  import inspect
  from absl import logging
  import keras_tuner
  from model_search import controller
  from model_search import hparam
  from model_search import loss_fns
  from model_search import metric_fns
  from model_search import task_manager
  from model_search.architecture import architecture_utils
  from model_search.generators import base_tower_generator
  from model_search.generators import trial_utils
  from model_search.meta import transfer_learning
  from model_search.metadata import ml_metadata_db
  from model_search.proto import phoenix_spec_pb2
  from model_search.proto import transfer_learning_spec_pb2
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 5 for calling the Phoenix obj's underlying member method code with example usage: obj._make_model_fn(run_config, train_steps, use_tpu) and returns: model_fn, tf, tf, self