| 1482 | } |
| 1483 | |
| 1484 | XlaOp BesselI0e(XlaOp x) { |
| 1485 | auto& b = *x.builder(); |
| 1486 | return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1487 | TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("BesselI0e", x)); |
| 1488 | TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(x)); |
| 1489 | if (shape.element_type() == F64) { |
| 1490 | return I0eImpl64(x); |
| 1491 | } |
| 1492 | // I0eF32Impl don't have enough precision when run with bf16 intermediates |
| 1493 | // (not surprising!), so upcast to f32 in this case. |
| 1494 | return DoWithUpcastToF32(x, {BF16, F16}, |
| 1495 | [](XlaOp x) { return I0eImpl32(x); }); |
| 1496 | }); |
| 1497 | } |
| 1498 | |
| 1499 | // Computes an approximation to the modified Bessel function of the first kind, |
| 1500 | // first order. |