| 216 | } |
| 217 | |
| 218 | StatusOr<std::vector<XlaOp>> WhileLoopFn( |
| 219 | absl::Span<const XlaOp> initial_values, // |
| 220 | int matrix_dimension, // |
| 221 | int max_sweep_updates, // |
| 222 | PrimitiveType index_type, // |
| 223 | absl::string_view name, // |
| 224 | XlaBuilder* builder) { |
| 225 | auto while_cond_fn = [&](absl::Span<const XlaOp> values, |
| 226 | XlaBuilder* cond_builder) -> StatusOr<XlaOp> { |
| 227 | auto k = values[0]; |
| 228 | auto max_sweeps = ScalarLike(k, max_sweep_updates); |
| 229 | auto sweep_update_cond = Gt(max_sweeps, k); |
| 230 | |
| 231 | auto norms = ComputeFrobeniusNorms(values[2]).ValueOrDie(); |
| 232 | auto tol = norms.total_norm * values[3]; |
| 233 | auto tol_cond = ReduceAll(Lt(tol, norms.off_diagonal_norm), |
| 234 | xla::ConstantR0<bool>(cond_builder, false), |
| 235 | CreateScalarOrComputation(PRED, cond_builder)); |
| 236 | |
| 237 | return And(sweep_update_cond, tol_cond); |
| 238 | }; |
| 239 | |
| 240 | auto while_body_fn = |
| 241 | [&](absl::Span<const XlaOp> values, |
| 242 | XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { |
| 243 | auto while_cond_fn_inner = |
| 244 | [&](absl::Span<const XlaOp> values_inner, |
| 245 | XlaBuilder* inner_cond_builder) -> StatusOr<XlaOp> { |
| 246 | auto p = values_inner[0]; |
| 247 | return Lt(p, ScalarLike(p, matrix_dimension - 1)); |
| 248 | }; |
| 249 | |
| 250 | auto while_body_fn_inner = |
| 251 | [&](absl::Span<const XlaOp> values_inner, |
| 252 | XlaBuilder* inner_body_builder) -> StatusOr<std::vector<XlaOp>> { |
| 253 | auto while_cond_fn_innermost = |
| 254 | [&](absl::Span<const XlaOp> values_innermost, |
| 255 | XlaBuilder* innermost_cond_builder) -> StatusOr<XlaOp> { |
| 256 | auto q = values_innermost[1]; |
| 257 | return Lt(q, ScalarLike(q, matrix_dimension)); |
| 258 | }; |
| 259 | auto while_body_fn_innermost = |
| 260 | [&](absl::Span<const XlaOp> values_innermost, |
| 261 | XlaBuilder* innermost_body_builder) |
| 262 | -> StatusOr<std::vector<XlaOp>> { |
| 263 | auto p = values_innermost[0]; |
| 264 | auto q = values_innermost[1]; |
| 265 | |
| 266 | JacobiUpdate jacobi_update; |
| 267 | jacobi_update.v = values_innermost[2]; |
| 268 | jacobi_update.w = values_innermost[3]; |
| 269 | |
| 270 | auto tol = values_innermost[4]; |
| 271 | |
| 272 | TF_ASSIGN_OR_RETURN(jacobi_update, |
| 273 | Update(jacobi_update, p, q, tol, matrix_dimension)); |
| 274 | |
| 275 | std::vector<XlaOp> updated_values_innermost; |
no test coverage detected