Show expected shape on TF API errors

Author: deliahu

pkg/workloads/cortex/tf_api/api.py

@@ -136,10 +136,16 @@ def create_prediction_request(transformed_sample):
         shape = []
         for dim in signature_def[signature_key]["tensorShape"]["dim"]:
             shape.append(int(dim["size"]))
-        tensor_proto = tf.make_tensor_proto(
-            np.array(value).reshape(shape), dtype=data_type, shape=shape
-        )
-        prediction_request.inputs[column_name].CopyFrom(tensor_proto)
+
+        try:
+            tensor_proto = tf.make_tensor_proto(
+                np.array(value).reshape(shape), dtype=data_type, shape=shape
+            )
+            prediction_request.inputs[column_name].CopyFrom(tensor_proto)
+        except Exception as e:
+            raise UserException(
+                'key "{}"'.format(column_name), "expected shape {}".format(shape)
+            ) from e
 
     return prediction_request
 
@@ -163,8 +169,16 @@ def create_raw_prediction_request(sample):
             shape = [1]
             value = [value]
         sig_type = signature_def[signature_key]["inputs"][column_name]["dtype"]
-        tensor_proto = tf.make_tensor_proto(value, dtype=DTYPE_TO_TF_TYPE[sig_type], shape=shape)
-        prediction_request.inputs[column_name].CopyFrom(tensor_proto)
+
+        try:
+            tensor_proto = tf.make_tensor_proto(
+                value, dtype=DTYPE_TO_TF_TYPE[sig_type], shape=shape
+            )
+            prediction_request.inputs[column_name].CopyFrom(tensor_proto)
+        except Exception as e:
+            raise UserException(
+                'key "{}"'.format(column_name), "expected shape {}".format(shape)
+            ) from e
 
     return prediction_request