aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/compiler/xla/tests/test_utils.cc
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/compiler/xla/tests/test_utils.cc')
-rw-r--r--tensorflow/compiler/xla/tests/test_utils.cc6
1 files changed, 4 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc
index 5155f0c652..2f18036ff4 100644
--- a/tensorflow/compiler/xla/tests/test_utils.cc
+++ b/tensorflow/compiler/xla/tests/test_utils.cc
@@ -272,9 +272,11 @@ std::vector<HloInstruction*> FindConstrainedUses(
constrained_uses.insert(constrained_uses.end(), converted_uses.begin(),
converted_uses.end());
} else if (opcode == HloOpcode::kSort &&
- instruction->operand_count() == 2 && op_num == 0) {
+ instruction->operand_count() >= 2 && op_num == 0) {
// Operand 0 of sort is the array of keys used for key/value
- // (two-operand) kSort instructions.
+ // (two-operand) kSort instructions. Since sort stability is not
+ // guaranteed, constrain keys of key-value sort not to have duplicates,
+ // since otherwise the value order may legitimately differ.
constrained_uses.push_back(instruction);
}
}