message: Simplify the boosted trees control flow.
diff: @@ -703,9 +703,30 @@ def _bt_model_fn(
  global_step = training_util.get_or_create_global_step()
  bucket_size_list, feature_ids_list = _group_features_by_num_buckets(
  sorted_feature_columns)
  + # Create Ensemble resources.
  + tree_ensemble = boosted_trees_ops.TreeEnsemble(name=name)
  +
  + # Create logits.
  + if mode != model_fn.ModeKeys.TRAIN:
  + input_feature_list = _get_transformed_features(features,
  + sorted_feature_columns)
  + logits = boosted_trees_ops.predict(
  + # For non-TRAIN mode, ensemble doesn't change after initialization,
  + # so no local copy is needed; using tree_ensemble directly.
  + tree_ensemble_handle=tree_ensemble.resource_handle,
  + bucketized_features=input_feature_list,
  + logits_dimension=head.logits_dimension)
  + return head.create_estimator_spec(
  + features=features,
  + mode=mode,
  + labels=labels,
  + train_op_fn=control_flow_ops.no_op,
  + logits=logits)
  +
  + # ============== Training graph ==============
  # Extract input features and set up cache for training.
  training_state_cache = None
  - if mode == model_fn.ModeKeys.TRAIN and train_in_memory:
  + if train_in_memory:
  # cache transformed features as well for in-memory training.
  batch_size = array_ops.shape(labels)[0]
  input_feature_list, input_cache_op = (
  @@ -717,25 +738,14 @@ def _bt_model_fn(
  else:
  input_feature_list = _get_transformed_features(features,
  sorted_feature_columns)
  - if mode == model_fn.ModeKeys.TRAIN and example_id_column_name:
  + if example_id_column_name:
  example_ids = features[example_id_column_name]
  training_state_cache = _CacheTrainingStatesUsingHashTable(
  example_ids, head.logits_dimension)
  - # Create Ensemble resources.
  - tree_ensemble = boosted_trees_ops.TreeEnsemble(name=name)
  # Variable that determines whether bias centering is needed.
  center_bias_var = variable_scope.variable(
  initial_value=center_bias, name='center_bias_needed', trainable=False)
  - # Create logits.
  - if mode != model_fn.ModeKeys.TRAIN:
  - logits = boosted_trees_ops.predict(
  - # For non-TRAIN mode, ensemble doesn't change after initialization,
  - # so no local copy is needed; using tree_ensemble directly.
  - tree_ensemble_handle=tree_ensemble.resource_handle,
  - bucketized_features=input_feature_list,
  - logits_dimension=head.logits_dimension)
  - else:
  if is_single_machine:
  local_tree_ensemble = tree_ensemble
  ensemble_reload = control_flow_ops.no_op()
  @@ -773,7 +783,6 @@ def _bt_model_fn(
  cached_node_ids=cached_node_ids,
  bucketized_features=input_feature_list,
  logits_dimension=head.logits_dimension)
  -
  logits = cached_logits + partial_logits
  # Create training graph.
  @@ -846,7 +855,6 @@ def _bt_model_fn(
  labels=labels,
  train_op_fn=_train_op_fn,
  logits=logits)
  - if mode == model_fn.ModeKeys.TRAIN:
  # Add an early stop hook.
  estimator_spec = estimator_spec._replace(
  training_hooks=estimator_spec.training_hooks +
  