[XLA] Shape inference should verify the shapes of sort keys and sort values match.

PiperOrigin-RevId: 204974328

1 files changed, 12 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc
index 9b1ce143c6..6046d50c6d 100644
--- a/tensorflow/compiler/xla/service/shape_inference_test.cc
+++ b/tensorflow/compiler/xla/service/shape_inference_test.cc
@@ -1524,6 +1524,18 @@ TEST_F(ShapeInferenceTest, BadSlice) {
       << statusor.status();
 }
 
+TEST_F(ShapeInferenceTest, BadSort) {
+  auto keys = ShapeUtil::MakeShape(F32, {4});
+  auto values = ShapeUtil::MakeShape(F32, {5});
+  StatusOr<Shape> statusor =
+      ShapeInference::InferVariadicOpShape(HloOpcode::kSort, {&keys, &values});
+  ASSERT_FALSE(statusor.ok());
+
+  EXPECT_THAT(statusor.status().error_message(),
+              HasSubstr("dimensions must match"))
+      << statusor.status();
+}
+
 class GatherShapeInferenceTest : public ShapeInferenceTest {
  protected:
   const Shape s64_scalar_ = ShapeUtil::MakeShape(S64, {});