Fix SavedModel export when predictions is a single tensor and output_alternatives not given

Change: 144470271

tensorflow/contrib/learn/python/learn/estimators/prediction_key.py

@@ -25,3 +25,4 @@ class PredictionKey(object):
   LOGISTIC = "logistic"
   SCORES = "scores"
   TOP_K = "top_k"
+  GENERIC = "output"

tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py

@@ -167,6 +167,8 @@ def get_output_alternatives(
   # interpret the model as single-headed of unknown type.
   default_problem_type = constants.ProblemType.UNSPECIFIED
   default_outputs = model_fn_ops.predictions
+  if not isinstance(default_outputs, dict):
+    default_outputs = {prediction_key.PredictionKey.GENERIC: default_outputs}
   actual_default_output_alternative_key = DEFAULT_OUTPUT_ALTERNATIVE_KEY
   output_alternatives = {actual_default_output_alternative_key:
                          (default_problem_type, default_outputs)}

tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py

@@ -124,6 +124,21 @@ class SavedModelExportUtilsTest(test.TestCase):
         })
     }, output_alternatives)
 
+  def test_get_output_alternatives_implicit_single(self):
+    prediction_tensor = constant_op.constant(["bogus"])
+    model_fn_ops = model_fn.ModelFnOps(
+        model_fn.ModeKeys.INFER,
+        predictions=prediction_tensor,
+        output_alternatives=None)
+
+    output_alternatives, _ = saved_model_export_utils.get_output_alternatives(
+        model_fn_ops)
+    self.assertEqual({
+        "default_output_alternative": (constants.ProblemType.UNSPECIFIED, {
+            "output": prediction_tensor
+        })
+    }, output_alternatives)
+
   def test_build_all_signature_defs(self):
     input_features = constant_op.constant(["10"])
     input_example = constant_op.constant(["11"])