Function IgammaSeries

tensorflow/compiler/xla/client/lib/math.cc:700–771

// Helper function for computing Igamma using a power series.
template <kIgammaMode mode>
XlaOp IgammaSeries(XlaOp ax, XlaOp x, XlaOp a, XlaOp enabled,
                   xla::PrimitiveType type) {
  // vals: (enabled, r, c, ans, x)
  // 'enabled' is a predication mask that says for which elements we should
  // execute the loop body. Disabled elements have no effect in the loop body.
  // TODO(phawkins): in general this isn't an optimal implementation on any
  // backend. For example, on GPU, we should probably vectorize to the warp
  // size, and then run independent loops for each warp's worth of
  // data.
  auto cond = [&](absl::Span<const XlaOp> vals,
                  XlaBuilder* builder) -> StatusOr<XlaOp> {
    XlaOp enabled = vals[0];
    return Any(enabled);
  };
  auto body = [&](absl::Span<const XlaOp> vals,
                  XlaBuilder* builder) -> StatusOr<std::vector<XlaOp>> {
    XlaOp enabled = vals[0];
    XlaOp r = vals[1];
    XlaOp c = vals[2];
    XlaOp ans = vals[3];
    XlaOp x = vals[4];
    XlaOp dc_da = vals[5];
    XlaOp dans_da = vals[6];

    r = r + ScalarLike(r, 1);
    dc_da = dc_da * (x / r) + (ScalarLike(r, -1) * c * x) / (r * r);
    dans_da = dans_da + dc_da;
    c = c * (x / r);
    ans = ans + c;
    XlaOp conditional;
    if (mode == VALUE) {
      conditional = And(enabled, Gt(c / ans, Epsilon(builder, type)));
    } else {
      conditional =
          And(enabled, Gt(Abs(dc_da / dans_da), Epsilon(builder, type)));
    }

    return std::vector<XlaOp>{
        conditional,
        Select(enabled, r, vals[1]),
        Select(enabled, c, vals[2]),
        Select(enabled, ans, vals[3]),
        Select(enabled, x, vals[4]),
        Select(enabled, dc_da, vals[5]),
        Select(enabled, dans_da, vals[6]),
    };
  };
  auto& b = *ax.builder();
  return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
    std::vector<XlaOp> vals = {
        enabled,        a, FullLike(a, 1), FullLike(a, 1), x, FullLike(a, 0),
        FullLike(a, 0),
    };

    TF_ASSIGN_OR_RETURN(vals, WhileLoopHelper(cond, body, vals, "igamma", &b));
    XlaOp ans = vals[3];
    XlaOp dans_da = vals[6];
    if (mode == VALUE) {