Switch FakeQuant and Dequantize ops in array_ops to use C++ shape functions.

Change: 137756007

Author: tensorflower-gardener

tensorflow/core/ops/array_ops.cc

@@ -4392,6 +4392,7 @@ REGISTER_OP("FakeQuantWithMinMaxArgs")
     .Attr("max: float = 6.0")
     .Input("inputs: float")
     .Output("outputs: float")
+    .SetShapeFn(shape_inference::UnchangedShape)
     .Doc(R"doc(
 Fake-quantize the 'inputs' tensor, type float to 'outputs' tensor of same type.
 
@@ -4422,6 +4423,13 @@ REGISTER_OP("FakeQuantWithMinMaxVars")
     .Input("min: float")
     .Input("max: float")
     .Output("outputs: float")
+    .SetShapeFn([](InferenceContext* c) {
+      TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
+      ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
+      return Status::OK();
+    })
     .Doc(R"doc(
 Fake-quantize the 'inputs' tensor of type float and shape `[b, h, w, d]` via
 global float scalars `min` and `max` to 'outputs' tensor of same shape as
@@ -4461,6 +4469,20 @@ REGISTER_OP("FakeQuantWithMinMaxVarsPerChannel")
     .Input("min: float")
     .Input("max: float")
     .Output("outputs: float")
+    .SetShapeFn([](InferenceContext* c) {
+      ShapeHandle input, min, max;
+      TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &input));
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &min));
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &max));
+
+      DimensionHandle unused;
+      TF_RETURN_IF_ERROR(c->Merge(c->Dim(input, -1), c->Dim(min, 0), &unused));
+      TF_RETURN_IF_ERROR(c->Merge(c->Dim(input, -1), c->Dim(max, 0), &unused));
+      TF_RETURN_IF_ERROR(c->Merge(c->Dim(min, 0), c->Dim(max, 0), &unused));
+
+      c->set_output(0, input);
+      return Status::OK();
+    })
     .Doc(R"doc(
 Fake-quantize the 'inputs' tensor of type float and one of the shapes: `[d]`,
 `[b, d]` `[b, h, w, d]` via per-channel floats `min` and `max` of shape `[d]`

tensorflow/core/ops/array_ops_test.cc

@@ -1493,4 +1493,36 @@ TEST(ArrayOpsTest, StridedSliceGrad_ShapeFn) {
   INFER_OK(op, "[4];?;?;?;?", "[1,2,3,4]");
 }
 
+TEST(ArrayOpsTest, UnchangedWithQuantizationScalars_ShapeFn) {
+  for (const char* op_name : {"Dequantize", "FakeQuantWithMinMaxVars"}) {
+    ShapeInferenceTestOp op(op_name);
+
+    INFER_OK(op, "?;?;?", "in0");
+    INFER_OK(op, "[1,?,3];[];[]", "in0");
+
+    // Rank check scalars.
+    INFER_ERROR("be rank 0", op, "[1,?,3];[1];[]");
+    INFER_ERROR("be rank 0", op, "[1,?,3];[];[1]");
+  }
+}
+
+TEST(ArrayOpsTest, FakeQuantWithMinMaxVarsPerChannel) {
+  ShapeInferenceTestOp op("FakeQuantWithMinMaxVarsPerChannel");
+
+  INFER_OK(op, "?;?;?", "?");
+  INFER_OK(op, "[?];?;?", "in0");
+  INFER_OK(op, "[1,?,3];[3];[3]", "in0");
+  INFER_OK(op, "[3];[3];[3]", "in0");
+
+  // Rank check vectors.
+  INFER_ERROR("be rank 1", op, "[1,?,3];[1];[]");
+  INFER_ERROR("be rank 1", op, "[1,?,3];[];[1]");
+
+  // Vectors must match each other, and match last dim of input.
+  INFER_ERROR("must be equal", op, "[1,?,3];[2];[?]");
+  INFER_ERROR("must be equal", op, "[1,?,3];[?];[2]");
+  INFER_ERROR("must be equal", op, "[1,?,?];[1];[2]");
+  INFER_ERROR("must be equal", op, "[5];[4];[?]");
+}
+
 }  // end namespace tensorflow

tensorflow/python/ops/array_ops.py

@@ -2011,18 +2011,8 @@ def _EditDistanceShape(op):
   return common_shapes.call_cpp_shape_fn(op, input_tensors_needed=[2, 5])
 
 
-@ops.RegisterShape("Quantize")
-@ops.RegisterShape("Dequantize")
-def _QuantizeDequantizeShape(op):
-  unused_min_range = op.inputs[1].get_shape().merge_with(tensor_shape.scalar())
-  unused_max_range = op.inputs[2].get_shape().merge_with(tensor_shape.scalar())
-  return common_shapes.unchanged_shape(op)
-
-
-@ops.RegisterShape("FakeQuantWithMinMaxArgs")
-def _FakeQuantWithMinMaxArgsShape(op):
-  """Shape function for FakeQuantWithMinMaxArgs op: preserve the input shape."""
-  return [op.inputs[0].get_shape()]
+ops.RegisterShape("Quantize")(common_shapes.call_cpp_shape_fn)
+ops.RegisterShape("Dequantize")(common_shapes.call_cpp_shape_fn)
 
 
 @ops.RegisterGradient("FakeQuantWithMinMaxArgs")
@@ -2031,10 +2021,10 @@ def _FakeQuantWithMinMaxArgsGradient(op, grad):
   return fake_quant_with_min_max_args_gradient(grad, op.inputs[0])
 
 
-@ops.RegisterShape("FakeQuantWithMinMaxVars")
-def _FakeQuantWithMinMaxVarsShape(op):
-  """Shape function for FakeQuantWithMinMaxVars op: preserve the input shape."""
-  return [op.inputs[0].get_shape()]
+ops.RegisterShape("FakeQuantWithMinMaxArgs")(common_shapes.call_cpp_shape_fn)
+ops.RegisterShape("FakeQuantWithMinMaxVars")(common_shapes.call_cpp_shape_fn)
+ops.RegisterShape("FakeQuantWithMinMaxVarsPerChannel")(
+    common_shapes.call_cpp_shape_fn)
 
 
 @ops.RegisterGradient("FakeQuantWithMinMaxVars")
@@ -2044,12 +2034,6 @@ def _FakeQuantWithMinMaxVarsGradient(op, grad):
                                                op.inputs[2])
 
 
-@ops.RegisterShape("FakeQuantWithMinMaxVarsPerChannel")
-def _FakeQuantWithMinMaxVarsPerChannelShape(op):
-  """Shape function for FakeQuantWithMinMaxVarsPerChannel op: input shape."""
-  return [op.inputs[0].get_shape()]
-
-
 @ops.RegisterGradient("FakeQuantWithMinMaxVarsPerChannel")
 def _FakeQuantWithMinMaxVarsPerChannelGradient(op, grad):
   """Gradient for FakeQuantWithMinMaxVarsPerChannel op."""