From a12b8c4afdca3ac2945d62b3b83ca2599ab360f9 Mon Sep 17 00:00:00 2001 From: Keno Fischer Date: Sun, 16 Sep 2018 18:39:50 -0400 Subject: [xla] Improve validation of Broadcast shape If one misreads the semantics of this instruction, it's easy to cause an out of bounds access into the dimensions here. Add an extra check to return a proper error to the user rather than crashing in that case. Ref #22130 --- tensorflow/compiler/xla/service/hlo_verifier.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'tensorflow/compiler') diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 50f39cbcb5..0f6ecd42f6 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -313,8 +313,9 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { operand_dimension < ShapeUtil::Rank(operand_shape); ++operand_dimension) { int64 output_dimension = broadcast->dimensions()[operand_dimension]; - TF_RET_CHECK(broadcast->shape().dimensions(output_dimension) == - operand_shape.dimensions(operand_dimension)) + TF_RET_CHECK((output_dimension < ShapeUtil::Rank(broadcast->shape())) && + (broadcast->shape().dimensions(output_dimension) == + operand_shape.dimensions(operand_dimension))) << broadcast->ToString() << " operand shape " << operand_shape; } return Status::OK(); -- cgit v1.2.3