From 65e79cd8bf32ca5d58b5b7fb0b344fb13bb794a0 Mon Sep 17 00:00:00 2001
From: Alexander Pivovarov <pivovaa@amazon.com>
Date: Tue, 25 Jun 2024 22:10:59 +0000
Subject: [PATCH] Add kPower case to IsNonNegative

---
 xla/service/algebraic_simplifier.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/xla/service/algebraic_simplifier.cc b/xla/service/algebraic_simplifier.cc
index dd13ffebd9e852..3b51133b7a83fa 100644
--- a/xla/service/algebraic_simplifier.cc
+++ b/xla/service/algebraic_simplifier.cc
@@ -229,6 +229,9 @@ bool IsNonNegative(const HloInstruction* hlo,
       return IsNonNegative(hlo->operand(0), options) ||
              IsNonNegative(hlo->operand(1), options);
     }
+    case HloOpcode::kPower: {
+      return IsNonNegative(hlo->operand(0), options);
+    }
     case HloOpcode::kSelect: {
       return IsNonNegative(hlo->operand(1), options) &&
              IsNonNegative(hlo->operand(2), options);