| 135 | } |
| 136 | |
| 137 | XlaOp IsNegZero(XlaOp operand) { |
| 138 | auto& b = *operand.builder(); |
| 139 | return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 140 | TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsNegZero", operand)); |
| 141 | TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); |
| 142 | |
| 143 | // The bitwise representation of -0 in bfloat16 and IEEE 754 is 0x80...0 |
| 144 | // (sign bit on, all other bits off). |
| 145 | switch (shape.element_type()) { |
| 146 | case F64: |
| 147 | return Eq(BitcastConvertType(operand, U64), |
| 148 | ConstantR0WithType(&b, U64, uint64{1} << 63)); |
| 149 | case F32: |
| 150 | return Eq(BitcastConvertType(operand, U32), |
| 151 | ConstantR0WithType(&b, U32, uint32{1} << 31)); |
| 152 | case F16: |
| 153 | case BF16: |
| 154 | // Not all XLA backends handle U16 well, so we convert to F32/U32. |
| 155 | // TODO(jlebar): It would be nice if we could stay in (B)F16/U16 for |
| 156 | // backends that *do* support it. |
| 157 | return Eq(BitcastConvertType(ConvertElementType(operand, F32), U32), |
| 158 | ConstantR0WithType(&b, U32, uint32{1} << 31)); |
| 159 | default: |
| 160 | LOG(FATAL) << "Expected real fp type."; |
| 161 | } |
| 162 | }); |
| 163 | } |
| 164 | |
| 165 | XlaOp Square(XlaOp operand) { return operand * operand; } |
| 166 | |