message: Add leaf index modes as an argument, which specifies when output leaf indices.
diff: @@ -287,7 +287,8 @@ class GradientBoostedDecisionTreeModel(object):
  loss_reduction=losses.Reduction.SUM_OVER_NONZERO_WEIGHTS,
  feature_columns=None,
  use_core_columns=False,
  - output_leaf_index=False):
  + output_leaf_index=False,
  + output_leaf_index_modes=None):
  """Construct a new GradientBoostedDecisionTreeModel function.
  Args:
  @@ -307,6 +308,9 @@ class GradientBoostedDecisionTreeModel(object):
  used.
  output_leaf_index: A boolean variable indicating whether to output leaf
  index into predictions dictionary.
  + output_leaf_index_modes: A list of modes from (TRAIN, EVAL, INFER) which
  + dictates when leaf indices will be outputted. By default, leaf indices
  + are only outputted in INFER mode.
  Raises:
  ValueError: if inputs are not valid.
  @@ -404,7 +408,16 @@ class GradientBoostedDecisionTreeModel(object):
  self._learner_config.multi_class_strategy ==
  learner_pb2.LearnerConfig.TREE_PER_CLASS and
  learner_config.num_classes == 2)
  +
  + if output_leaf_index_modes is None:
  + output_leaf_index_modes = [learn.ModeKeys.INFER]
  + elif not all(
  + mode in (learn.ModeKeys.TRAIN, learn.ModeKeys.EVAL,
  + learn.ModeKeys.INFER) for mode in output_leaf_index_modes):
  + raise ValueError("output_leaf_index_modes should only contain ModeKeys.")
  +
  self._output_leaf_index = output_leaf_index
  + self._output_leaf_index_modes = output_leaf_index_modes
  def _predict_and_return_dict(self, ensemble_handle, ensemble_stamp, mode):
  """Runs prediction and returns a dictionary of the prediction results.
  @@ -435,8 +448,7 @@ class GradientBoostedDecisionTreeModel(object):
  # the right stamp.
  with ops.control_dependencies(ensemble_stats):
  leaf_index = None
  - # Only used in infer (predict), not used in train and eval.
  - if self._output_leaf_index and mode == learn.ModeKeys.INFER:
  + if self._output_leaf_index and mode in self._output_leaf_index_modes:
  predictions, _, leaf_index = (
  prediction_ops).gradient_trees_prediction_verbose(
  ensemble_handle,
  