from tensorflow.contrib.learn.python.learn.estimators import tensor_signature


  def predict_proba(self, x, batch_size=None):
    """Returns prediction probabilities for given features (classification).

    Args:
      x: features.
      batch_size: OVerride default batch size.

    Returns:
      Numpy array of predicted probabilities.
    """
    return self._infer_model(x=x, batch_size=batch_size, proba=True)

  def _check_inputs(self, features, targets):
    if self._features_info is not None:
      if not tensor_signature.tensors_compatible(features, self._features_info):
        raise ValueError('Features are incompatible with given information. '
                         'Given features: %s, required signatures: %s.' %
                         (str(features), str(self._features_info)))
    else:
      self._features_info = tensor_signature.create_signatures(features)
    if self._targets_info is not None:
      if not tensor_signature.tensors_compatible(targets, self._targets_info):
        raise ValueError('Targets are incompatible with given information. '
                         'Given targets: %s, required signatures: %s.' %
                         (str(targets), str(self._targets_info)))
    else:
      self._targets_info = tensor_signature.create_signatures(targets)

  def _train_model(self,
