| 2344 | } |
| 2345 | |
| 2346 | Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) { |
| 2347 | // ln(exp(A)) => A |
| 2348 | VLOG(10) << "trying transform [ln(exp(A)) => A]: " << log->ToString(); |
| 2349 | HloInstruction *a, *b; |
| 2350 | if (Match(log, m::Log(m::Exp(m::Op(&a)))) && |
| 2351 | ReplaceInstructionIfSameShape(log, a)) { |
| 2352 | return Status::OK(); |
| 2353 | } |
| 2354 | |
| 2355 | // ln(pow(A,B)) => B*ln(abs(A)) |
| 2356 | // or B*ln(A) if A is complex. |
| 2357 | if (Match(log, m::Log(m::Power(m::Op(&a), m::Op(&b))))) { |
| 2358 | auto abs_a = ShapeUtil::ElementIsComplex(a->shape()) |
| 2359 | ? a |
| 2360 | : computation_->AddInstruction(HloInstruction::CreateUnary( |
| 2361 | log->shape(), HloOpcode::kAbs, a)); |
| 2362 | auto new_log = computation_->AddInstruction( |
| 2363 | HloInstruction::CreateUnary(log->shape(), HloOpcode::kLog, abs_a)); |
| 2364 | return ReplaceWithNewInstruction( |
| 2365 | log, HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, |
| 2366 | new_log, b)); |
| 2367 | } |
| 2368 | |
| 2369 | if (Match(log, m::Log(m::Sqrt(m::Op(&a))))) { |
| 2370 | auto new_log = computation_->AddInstruction( |
| 2371 | HloInstruction::CreateUnary(log->shape(), HloOpcode::kLog, a)); |
| 2372 | return ReplaceWithNewInstruction( |
| 2373 | log, HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, |
| 2374 | new_log, MakeScalarLike(log, 0.5))); |
| 2375 | } |
| 2376 | |
| 2377 | if (Match(log, m::Log(m::Rsqrt(m::Op(&a))))) { |
| 2378 | auto new_log = computation_->AddInstruction( |
| 2379 | HloInstruction::CreateUnary(log->shape(), HloOpcode::kLog, a)); |
| 2380 | return ReplaceWithNewInstruction( |
| 2381 | log, HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, |
| 2382 | new_log, MakeScalarLike(log, -0.5))); |
| 2383 | } |
| 2384 | |
| 2385 | return Status::OK(); |
| 2386 | } |
| 2387 | |
| 2388 | Status AlgebraicSimplifierVisitor::HandleGetTupleElement( |
| 2389 | HloInstruction* get_tuple_element) { |