from tensorflow.python.ops import math_ops

    logits = self._logits(features)
    return self._logits_to_predictions(logits, proba=True)

  def _logits_to_predictions(self, logits, proba=False):
    if self._n_classes < 2:
      return array_ops.reshape(logits, [-1])

    if self._n_classes == 2:
      logits = array_ops.concat(1, [array_ops.zeros_like(logits), logits])

    if proba:
      return nn.softmax(logits)
    else:
      return math_ops.argmax(logits, 1)

  def _get_feature_ops_from_example(self, examples_batch):
    column_types = layers.create_dict_for_parse_example(
        (self._get_linear_feature_columns() or []) +
        (self._get_dnn_feature_columns() or []))
    features = parsing_ops.parse_example(examples_batch, column_types)
    return features

  def _num_label_columns(self):
    return 1 if self._n_classes <= 2 else self._n_classes

  def _get_linear_feature_columns(self):
