Function RandomGammaGrad

tensorflow/compiler/xla/client/lib/math.cc:1002–1047

Gradient of Gamma sample from Gamma(a, 1) with respect to `a`.

// Gradient of Gamma sample from Gamma(a, 1) with respect to `a`.
XlaOp RandomGammaGrad(XlaOp a, XlaOp x) {
  auto& b = *a.builder();
  auto doit = [&b](XlaOp a, XlaOp x, PrimitiveType type) -> XlaOp {
    XlaOp is_nan = Or(IsNan(a), IsNan(x));
    XlaOp x_is_zero = Eq(x, ScalarLike(x, 0));
    XlaOp domain_error = Or(Lt(x, ScalarLike(x, 0)), Le(a, ScalarLike(a, 0)));
    XlaOp use_igammac = And(Gt(x, ScalarLike(x, 1)), Gt(x, a));
    XlaOp ax = a * Log(x) - x - Lgamma(a);
    XlaOp underflow = Lt(ax, -Log(MaxFiniteValue(&b, type)));
    ax = Exp(ax);
    XlaOp enabled = Not(Or(Or(Or(x_is_zero, domain_error), underflow), is_nan));
    const double nan = std::numeric_limits<double>::quiet_NaN();
    XlaOp output = Select(use_igammac,
                          -IgammacContinuedFraction<SAMPLE_DERIVATIVE>(
                              ax, x, a, And(enabled, use_igammac), type),
                          IgammaSeries<SAMPLE_DERIVATIVE>(
                              ax, x, a, And(enabled, Not(use_igammac)), type));
    output = Select(underflow, ZerosLike(output), output);
    output = Select(x_is_zero, ZerosLike(output), output);
    output = Select(Or(domain_error, is_nan), FullLike(a, nan), output);
    return output;
  };
  return b.ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
    TF_ASSIGN_OR_RETURN(auto a_shape, b.GetShape(a));
    TF_ASSIGN_OR_RETURN(auto x_shape, b.GetShape(x));
    if (a_shape != x_shape) {
      return InvalidArgument(
          "Arguments to RandomGammaGrad must have equal shapes and types; got "
          "%s and %s",
          a_shape.ToString(), x_shape.ToString());
    }
    TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("RandomGammaGrad", a));
    bool needs_upcast =
        a_shape.element_type() == F16 || a_shape.element_type() == BF16;

    if (needs_upcast) {
      a = ConvertElementType(a, F32);
      x = ConvertElementType(x, F32);
    }
    XlaOp result = doit(a, x, a_shape.element_type());
    if (needs_upcast) {
      result = ConvertElementType(result, a_shape.element_type());
    }
    return result;
  });
}

XlaOp Igammac(XlaOp a, XlaOp x) {
  auto& b = *a.builder();