Require same shape for x and y in shape function of ApproximateEqual (tensorflow#19878)

yongtang · drpngx · commit 400a398a1878 · 2018-06-12T14:40:54.000-07:00
* Require same shape for `x` and `y` in shape function of `ApproximateEqual`

In the kernel implementation of `ApproximateEqual` the shape of inputs
`x` and `y` should be the same. Though in the shape function of `ApproximateEqual`
there was no such validation. This fix adds the shape validation in the
shape function to make sure `x` and `y` are of the same shape, if they are known.

Signed-off-by: Yong Tang &lt;yong.tang.github@outlook.com&gt;

* Add test case for shape function of ApproximateEqual

Signed-off-by: Yong Tang &lt;yong.tang.github@outlook.com&gt;
diff --git a/tensorflow/core/ops/math_ops.cc b/tensorflow/core/ops/math_ops.cc
@@ -592,7 +592,13 @@ REGISTER_OP("ApproximateEqual")
     .SetIsCommutative()
     .Attr("T: numbertype")
     .Attr("tolerance: float = 0.00001")
-    .SetShapeFn(shape_inference::UnchangedShape);
+    .SetShapeFn([](InferenceContext* c) {
+      // The inputs 'x' and 'y' must have the same shape.
+      ShapeHandle data_x = c->input(0);
+      ShapeHandle data_y = c->input(1);
+      TF_RETURN_IF_ERROR(c->Merge(data_x, data_y, &data_x));
+      return shape_inference::UnchangedShape(c);
+    });
 
 // --------------------------------------------------------------------------
 
diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
@@ -235,6 +235,15 @@ def testApproximateEqual(self):
         z_tf = self.evaluate(math_ops.approximate_equal(x, y, tolerance=0.0001))
         self.assertAllEqual(z, z_tf)
 
+  def testApproximateEqualShape(self):
+    for dtype in [np.float32, np.double]:
+      x = np.array([1, 2], dtype=np.float32)
+      y = np.array([[1, 2]], dtype=np.float32)
+      # The inputs 'x' and 'y' must have the same shape.
+      with self.assertRaisesRegexp(
+          ValueError, "Shapes must be equal rank, but are 1 and 2"):
+        math_ops.approximate_equal(x, y)
+
 
 class ScalarMulTest(test_util.TensorFlowTestCase):
 
