Function WhileLoopFn

tensorflow/compiler/xla/client/lib/svd.cc:588–714

Main boby of One-sided Jacobi Method.

// Main boby of One-sided Jacobi Method.
StatusOr<std::vector<XlaOp>> WhileLoopFn(
    absl::Span<const XlaOp> initial_values,  //
    int matrix_dimension,                    //
    int max_sweep_updates,                   //
    absl::string_view name,                  //
    XlaBuilder* builder) {
  auto while_cond_fn = [&](absl::Span<const XlaOp> values,
                           XlaBuilder* cond_builder) -> StatusOr<XlaOp> {
    auto k = values[0];
    auto max_sweeps = ScalarLike(k, max_sweep_updates);
    auto sweep_update_cond = Gt(max_sweeps, k);

    TF_ASSIGN_OR_RETURN(auto tolerance_comparison,
                        ComputeToleranceComparison(values[3], values[4]));
    auto tolerance_cond = ReduceAll(
        tolerance_comparison, xla::ConstantR0<bool>(cond_builder, false),
        CreateScalarOrComputation(PRED, cond_builder));

    return And(sweep_update_cond, tolerance_cond);
  };

  auto while_body_fn =
      [&](absl::Span<const XlaOp> values,
          XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> {
    auto while_cond_fn_inner =
        [&](absl::Span<const XlaOp> values_inner,
            XlaBuilder* inner_cond_builder) -> StatusOr<XlaOp> {
      auto p = values_inner[0];
      return Lt(p, ScalarLike(p, matrix_dimension - 1));
    };

    auto while_body_fn_inner =
        [&](absl::Span<const XlaOp> values_inner,
            XlaBuilder* inner_body_builder) -> StatusOr<std::vector<XlaOp>> {
      auto while_cond_fn_innermost =
          [&](absl::Span<const XlaOp> values_innermost,
              XlaBuilder* innermost_cond_builder) -> StatusOr<XlaOp> {
        auto q = values_innermost[1];
        return Lt(q, ScalarLike(q, matrix_dimension));
      };
      auto while_body_fn_innermost =
          [&](absl::Span<const XlaOp> values_innermost,
              XlaBuilder* innermost_body_builder)
          -> StatusOr<std::vector<XlaOp>> {
        auto p = values_innermost[0];
        auto q = values_innermost[1];

        SVDResult onesided_jacobi_update;
        onesided_jacobi_update.u = values_innermost[2];
        onesided_jacobi_update.v = values_innermost[3];
        onesided_jacobi_update.d = values_innermost[4];

        auto eps = values_innermost[5];

        TF_ASSIGN_OR_RETURN(
            onesided_jacobi_update,
            OneSidedJacobiUpdate(onesided_jacobi_update, p, q, eps));