From 6ddc7f8d99435c352b55aef79e7674973d6d3512 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 8 Jun 2020 22:26:52 -0700 Subject: [PATCH] [XLA] Update layout tiling in some algebraic simplifier cases. PiperOrigin-RevId: 315421103 Change-Id: Ic4dee365811f0faa5678d846a1e1412663610062 --- tensorflow/compiler/xla/service/algebraic_simplifier.cc | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index e0a8b87c83b..cd4dc05ab09 100755 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -2973,6 +2973,7 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { // slice instruction should all have the same layout. TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( pad->shape(), nonzero_pad->mutable_shape())); + simplifier_->UpdateLayout(nonzero_pad->mutable_shape()); // Second, construct the slice instruction to perform the negative padding. std::vector start_indices; @@ -2999,9 +3000,14 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { MakeSliceHlo(nonzero_pad, start_indices, end_indices, strides)); TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( pad->shape(), slice->mutable_shape())); + simplifier_->UpdateLayout(slice->mutable_shape()); // Verify that the slice shape matches the pad shape. - TF_RET_CHECK(ShapeUtil::Equal(slice->shape(), pad->shape())); + auto equal = Shape::Equal(); + if (!options_.is_layout_sensitive()) { + equal.IgnoreTilesInLayout(); + } + TF_RET_CHECK(equal(slice->shape(), pad->shape())); return ReplaceInstruction(pad, slice); }