| 1047 | } |
| 1048 | |
| 1049 | XlaOp Igammac(XlaOp a, XlaOp x) { |
| 1050 | auto& b = *a.builder(); |
| 1051 | auto doit = [&b](XlaOp a, XlaOp x, PrimitiveType type) -> XlaOp { |
| 1052 | XlaOp out_of_range = Or(Le(x, ScalarLike(x, 0)), Le(a, ScalarLike(a, 0))); |
| 1053 | XlaOp use_igamma = Or(Lt(x, ScalarLike(x, 1)), Lt(x, a)); |
| 1054 | XlaOp ax = a * Log(x) - x - Lgamma(a); |
| 1055 | XlaOp underflow = Lt(ax, -Log(MaxFiniteValue(&b, type))); |
| 1056 | XlaOp enabled = Not(Or(out_of_range, underflow)); |
| 1057 | ax = Exp(ax); |
| 1058 | XlaOp result = |
| 1059 | Select(use_igamma, |
| 1060 | ScalarLike(a, 1) - IgammaSeries<VALUE>( |
| 1061 | ax, x, a, And(enabled, use_igamma), type), |
| 1062 | IgammacContinuedFraction<VALUE>( |
| 1063 | ax, x, a, And(enabled, Not(use_igamma)), type)); |
| 1064 | return Select(underflow, ZerosLike(a), |
| 1065 | Select(out_of_range, FullLike(a, 1), result)); |
| 1066 | }; |
| 1067 | return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1068 | TF_ASSIGN_OR_RETURN(auto a_shape, b.GetShape(a)); |
| 1069 | TF_ASSIGN_OR_RETURN(auto x_shape, b.GetShape(x)); |
| 1070 | if (a_shape != x_shape) { |
| 1071 | return InvalidArgument( |
| 1072 | "Arguments to Igammac must have equal shapes and types; " |
| 1073 | "got %s and %s", |
| 1074 | a_shape.ToString(), x_shape.ToString()); |
| 1075 | } |
| 1076 | TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Igammac", a)); |
| 1077 | bool needs_upcast = |
| 1078 | a_shape.element_type() == F16 || a_shape.element_type() == BF16; |
| 1079 | |
| 1080 | if (needs_upcast) { |
| 1081 | a = ConvertElementType(a, F32); |
| 1082 | x = ConvertElementType(x, F32); |
| 1083 | } |
| 1084 | XlaOp result = doit(a, x, a_shape.element_type()); |
| 1085 | if (needs_upcast) { |
| 1086 | result = ConvertElementType(result, a_shape.element_type()); |
| 1087 | } |
| 1088 | return result; |
| 1089 | }); |
| 1090 | } |
| 1091 | // Implements Banker's rounding: numbers that are equidistant between two |
| 1092 | // integers are rounded towards even. |
| 1093 | XlaOp RoundToEven(XlaOp x) { |