[XLA] Don't simplify pow(pow(a,x),y) to pow(a,x*y) since x*y can not represent rational numbers.

PiperOrigin-RevId: 307560597
Change-Id: I89ce4ab52ad62386f6c6cb66b72a975d174693a1
This commit is contained in:
Blake Hechtman 2020-04-21 01:23:36 -07:00 committed by TensorFlower Gardener
parent 0ead27a263
commit 2d79e9922f
2 changed files with 1 additions and 26 deletions

View File

@ -2969,26 +2969,6 @@ Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) {
MakeScalarLike(lhs, 1), lhs));
}
VLOG(10) << "trying transform [pow(pow(A, X), Y) => pow(A, X*Y)]: "
<< power->ToString();
// Don't perform this optimization if either of the exponents is complex; this
// identity is true only for real-valued exponents. In addition, we cowardly
// refuse to do this transformation if the two exponents have different
// element types.
if (lhs->opcode() == HloOpcode::kPower &&
!ShapeUtil::ElementIsComplex(lhs->operand(1)->shape()) &&
!ShapeUtil::ElementIsComplex(rhs->shape()) &&
ShapeUtil::SameElementType(lhs->operand(1)->shape(), rhs->shape())) {
auto exponent_product =
computation_->AddInstruction(HloInstruction::CreateBinary(
rhs->shape(), HloOpcode::kMultiply, lhs->mutable_operand(1), rhs));
return ReplaceWithNewInstruction(
power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kPower,
lhs->mutable_operand(0),
exponent_product));
}
return Status::OK();
}

View File

@ -1011,13 +1011,8 @@ TEST_F(AlgebraicSimplifierTest, PowerOfPower) {
builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kPower,
inner_power, exp2));
auto computation = m->AddEntryComputation(builder.Build());
AlgebraicSimplifier simplifier(default_options_);
ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie());
EXPECT_THAT(
computation->root_instruction(),
GmockMatch(m::Power(m::Op().Is(base),
m::Multiply(m::Op().Is(exp1), m::Op().Is(exp2)))));
ASSERT_FALSE(simplifier.Run(m.get()).ValueOrDie());
}
// Don't simplify pow(pow(A, X), Y) => pow(A, X*Y) if X and Y are complex