| 906 | } // namespace |
| 907 | |
| 908 | XlaOp Igamma(XlaOp a, XlaOp x) { |
| 909 | auto& b = *a.builder(); |
| 910 | auto doit = [&b](XlaOp a, XlaOp x, PrimitiveType type) -> XlaOp { |
| 911 | XlaOp is_nan = Or(IsNan(a), IsNan(x)); |
| 912 | XlaOp x_is_zero = Eq(x, ScalarLike(x, 0)); |
| 913 | XlaOp domain_error = Or(Lt(x, ScalarLike(x, 0)), Le(a, ScalarLike(a, 0))); |
| 914 | XlaOp use_igammac = And(Gt(x, ScalarLike(x, 1)), Gt(x, a)); |
| 915 | XlaOp ax = a * Log(x) - x - Lgamma(a); |
| 916 | XlaOp underflow = Lt(ax, -Log(MaxFiniteValue(&b, type))); |
| 917 | ax = Exp(ax); |
| 918 | XlaOp enabled = Not(Or(Or(Or(x_is_zero, domain_error), underflow), is_nan)); |
| 919 | const double nan = std::numeric_limits<double>::quiet_NaN(); |
| 920 | XlaOp output = Select( |
| 921 | use_igammac, |
| 922 | ScalarLike(a, 1) - IgammacContinuedFraction<VALUE>( |
| 923 | ax, x, a, And(enabled, use_igammac), type), |
| 924 | IgammaSeries<VALUE>(ax, x, a, And(enabled, Not(use_igammac)), type)); |
| 925 | output = Select(underflow, ZerosLike(output), output); |
| 926 | output = Select(x_is_zero, ZerosLike(output), output); |
| 927 | output = Select(Or(domain_error, is_nan), FullLike(a, nan), output); |
| 928 | return output; |
| 929 | }; |
| 930 | return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 931 | TF_ASSIGN_OR_RETURN(auto a_shape, b.GetShape(a)); |
| 932 | TF_ASSIGN_OR_RETURN(auto x_shape, b.GetShape(x)); |
| 933 | if (a_shape != x_shape) { |
| 934 | return InvalidArgument( |
| 935 | "Arguments to Igamma must have equal shapes and types; got %s and %s", |
| 936 | a_shape.ToString(), x_shape.ToString()); |
| 937 | } |
| 938 | TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Igamma", a)); |
| 939 | bool needs_upcast = |
| 940 | a_shape.element_type() == F16 || a_shape.element_type() == BF16; |
| 941 | |
| 942 | if (needs_upcast) { |
| 943 | a = ConvertElementType(a, F32); |
| 944 | x = ConvertElementType(x, F32); |
| 945 | } |
| 946 | XlaOp result = doit(a, x, a_shape.element_type()); |
| 947 | if (needs_upcast) { |
| 948 | result = ConvertElementType(result, a_shape.element_type()); |
| 949 | } |
| 950 | return result; |
| 951 | }); |
| 952 | } |
| 953 | |
| 954 | XlaOp IgammaGradA(XlaOp a, XlaOp x) { |
| 955 | auto& b = *a.builder(); |