Function RegularizedIncompleteBeta

tensorflow/compiler/xla/client/lib/math.cc:1700–1825

}

XlaOp RegularizedIncompleteBeta(XlaOp a, XlaOp b, XlaOp x) {
  auto& builder = *x.builder();
  return builder.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
    TF_ASSIGN_OR_RETURN(Shape shape, builder.GetShape(a));
    TF_ASSIGN_OR_RETURN(Shape b_shape, builder.GetShape(b));
    TF_ASSIGN_OR_RETURN(Shape x_shape, builder.GetShape(x));
    if (b_shape.element_type() != shape.element_type() ||
        x_shape.element_type() != shape.element_type()) {
      return InvalidArgument(
          "Operands to RegularizedIncompleteBeta must have identical types, "
          "got shapes %s, %s, and %s",
          shape.ToString(), b_shape.ToString(), x_shape.ToString());
    }
    if (!primitive_util::IsFloatingPointType(shape.element_type())) {
      return InvalidArgument(
          "Operands to RegularizedIncompleteBeta must be real-valued "
          "floating-point, but got %s",
          PrimitiveType_Name(shape.element_type()));
    }
    PrimitiveType element_type = shape.element_type();
    if (element_type == F16 || element_type == BF16) {
      element_type = F32;
      a = ConvertElementType(a, F32);
      b = ConvertElementType(b, F32);
      x = ConvertElementType(x, F32);
    }

    // The partial numerator for the incomplete beta function is given
    // here: http://dlmf.nist.gov/8.17.E23 Note that there is a special
    // case: the partial numerator for the first iteration is one.
    auto NthPartialBetaincNumerator =
        [&](XlaOp iteration, absl::Span<const XlaOp> inputs,
            XlaBuilder* builder) -> StatusOr<std::vector<XlaOp>> {
      auto a = inputs[0];
      auto b = inputs[1];
      auto x = inputs[2];
      auto iteration_bcast = Broadcast(iteration, shape.dimensions());
      auto iteration_is_even =
          Eq(iteration_bcast % FullLike(iteration_bcast, 2),
             FullLike(iteration_bcast, 0));
      auto iteration_is_one = Eq(iteration_bcast, FullLike(iteration_bcast, 1));
      auto iteration_minus_one = iteration_bcast - FullLike(iteration_bcast, 1);
      auto m = iteration_minus_one / FullLike(iteration_minus_one, 2);
      m = ConvertElementType(m, element_type);
      auto one = FullLike(a, 1.0);
      auto two = FullLike(a, 2.0);
      // Partial numerator terms.
      auto even_numerator =
          -(a + m) * (a + b + m) * x / ((a + two * m) * (a + two * m + one));
      auto odd_numerator =
          m * (b - m) * x / ((a + two * m - one) * (a + two * m));
      auto one_numerator = ScalarLike(x, 1.0);
      auto numerator = Select(iteration_is_even, even_numerator, odd_numerator);
      return std::vector<XlaOp>{
          Select(iteration_is_one, one_numerator, numerator)};
    };

    auto NthPartialBetaincDenominator =