Returns operation(operand), except if `operand` is one of the types in upcast_types, in which case first converts it to F32, and then converts the result down to the original type.
| 63 | // upcast_types, in which case first converts it to F32, and then converts the |
| 64 | // result down to the original type. |
| 65 | static XlaOp DoWithUpcastToF32(XlaOp operand, |
| 66 | absl::Span<const PrimitiveType> upcast_types, |
| 67 | const std::function<XlaOp(XlaOp)>& operation) { |
| 68 | auto& b = *operand.builder(); |
| 69 | return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 70 | TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); |
| 71 | PrimitiveType elem_ty = shape.element_type(); |
| 72 | bool needs_upcast = absl::c_linear_search(upcast_types, elem_ty); |
| 73 | |
| 74 | if (needs_upcast) { |
| 75 | operand = ConvertElementType(operand, F32); |
| 76 | } |
| 77 | XlaOp result = operation(operand); |
| 78 | if (needs_upcast) { |
| 79 | result = ConvertElementType(result, elem_ty); |
| 80 | } |
| 81 | return result; |
| 82 | }); |
| 83 | } |
| 84 | |
| 85 | // TODO(jlebar): Use this function in more places in this file to restrict the |
| 86 | // domain of other functions. |
no test coverage detected